Fix langchain warnings

Author: piyushroshan

services/chatbot/src/chatbot/agent_utils.py

@@ -1,6 +1,9 @@
 import json
 
+from langchain.agents import AgentState
+from langchain.agents.middleware.types import before_model
 from langchain_core.messages import ToolMessage
+from langgraph.runtime import Runtime
 
 from .config import Config
 
@@ -60,7 +63,8 @@ def truncate_by_length(content, max_length):
     return content[:max_length] + "\n... [TRUNCATED]"
 
 
-def truncate_tool_messages(state):
+@before_model(state_schema=AgentState)
+def truncate_tool_messages(state: AgentState, runtime: Runtime) -> AgentState:
     """
     Modify large tool messages to prevent exceeding model's token limits.
     Truncate to a length such that it keeps messages within your token limit.

services/chatbot/src/chatbot/langgraph_agent.py

@@ -59,8 +59,8 @@ async def build_langgraph_agent(api_key, model_name, user_jwt):
     agent_node = create_agent(
         model=llm,
         tools=tools,
-        prompt=system_prompt,
-        pre_model_hook=truncate_tool_messages,
+        system_prompt=system_prompt,
+        middleware=[truncate_tool_messages],
     )
     return agent_node
 

services/chatbot/src/chatbot/retriever_utils.py

@@ -1,8 +1,8 @@
 import chromadb
 from langchain_chroma import Chroma as ChromaClient
-from langchain_community.embeddings import OpenAIEmbeddings
 from langchain_core.documents import Document
 from langchain_core.tools import create_retriever_tool
+from langchain_openai import OpenAIEmbeddings
 
 from .config import Config