Update main.py

main.py

@@ -7,6 +7,7 @@ from typing import List
 from pathlib import Path
 from langchain_huggingface import HuggingFaceEndpoint
 from langchain.schema import StrOutputParser
+from langchain_anthropic import AnthropicLLM
 
 from langchain.agents import AgentExecutor
 from langchain.agents.agent_types import AgentType
@@ -26,6 +27,11 @@ async def LLMistral():
         repo_id=repo_id, max_new_tokens=5300, temperature=0.1, task="text2text-generation", streaming=True
     )
     return llm
+
+async def LLM():
+    os.environ['ANTHROPIC_API_KEY'] = os.environ['ANTHROPIC_API_KEY']
+    llm = AnthropicLLM(model='claude-2.1')
+    return llm
     
 @cl.set_chat_profiles
 async def chat_profile():
@@ -49,18 +55,16 @@ async def on_message(message: cl.Message):
     df = pd.read_csv(
         "./public/ExpeCFA_LP_CAA.csv"
     )
-    model = await LLMistral()
+    model = await LLM()
     
-    agent = create_pandas_dataframe_agent(model, df, verbose=True, allow_dangerous_code=True, handle_parsing_errors=True)
-    
-    #agent = create_csv_agent(
-    #    model,
-    #    "./public/ExpeCFA_LP_CAA.csv",
-    #    verbose=False,
-    #    allow_dangerous_code=True,
-    #    handle_parsing_errors=True,
-    #    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION
-    #)
+    agent = create_csv_agent(
+        model,
+        "./public/ExpeCFA_LP_CAA.csv",
+        verbose=False,
+        allow_dangerous_code=True,
+        handle_parsing_errors=True,
+        agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION
+    )
     
     msg = cl.Message(content="")