add claude to models

ur-whitelab · Aug 7, 2024 · 2e7483e · 2e7483e
1 parent 794c9b5
commit 2e7483e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/mdagent/utils/makellm.py b/mdagent/utils/makellm.py
@@ -33,6 +33,15 @@ def _make_llm(model, temp, verbose):
             streaming=True if verbose else False,
             callbacks=[StreamingStdOutCallbackHandler()] if verbose else None,
         )
+    elif model.startswith("claude"):
+        from langchain_anthropic import ChatAnthropic
+
+        llm = ChatAnthropic(
+            temperature=temp,
+            model_name=model,
+            streaming=True if verbose else False,
+            callbacks=[StreamingStdOutCallbackHandler()] if verbose else None,
+        )
     else:
         raise ValueError(f"Unrecognized or unsupported model name: {model}")
     return llm
diff --git a/setup.py b/setup.py
@@ -20,6 +20,7 @@
         "chromadb",
         "google-search-results",
         "langchain==0.2.12",
+        "langchain-anthropic",
         "langchain-chroma",
         "langchain-community",
         "langchain-fireworks==0.1.7",