From 2e7483ea94a79832ff0a858656b51a0a20331dfb Mon Sep 17 00:00:00 2001
From: qcampbel <qcampbe2@ur.rochester.edu>
Date: Wed, 7 Aug 2024 14:37:01 -0400
Subject: [PATCH] add claude to models

---
 mdagent/utils/makellm.py | 9 +++++++++
 setup.py                 | 1 +
 2 files changed, 10 insertions(+)

diff --git a/mdagent/utils/makellm.py b/mdagent/utils/makellm.py
index 24a1fb0e..374e7560 100644
--- 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
index 99763e17..e753aa6e 100644
--- 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",