Update main.py

main.py

@@ -1,3 +1,7 @@
+#DOCS 
+# https://langchain-ai.github.io/langgraph/reference/prebuilt/#langgraph.prebuilt.chat_agent_executor.create_react_agent
+
+
 import uuid
 from fastapi import FastAPI
 from fastapi.responses import StreamingResponse
@@ -196,19 +200,30 @@ def get_collection_files(collection_id: str, user_id: str) -> str:
         logger.error(f"Error getting collection files: {str(e)}")
         return f"Error getting files: {str(e)}"
 
-def format_for_model(state):
-    # Get collection_id and user_id from the state's configurable
-    config = state.get("configurable", {})
-    collection_id = config.get("collection_id")
-    user_id = config.get("user_id")
+def format_for_model(state: dict, config: Optional[RunnableConfig] = None) -> list[BaseMessage]:
+    """
+    Format the input state and config for the model.
+    
+    Args:
+        state: The current state dictionary containing messages
+        config: Optional RunnableConfig containing thread configuration
+        
+    Returns:
+        Formatted messages for the model
+    """
+    # Get collection_id and user_id from config instead of state
+    thread_config = config.get("configurable", {}) if config else {}
+    collection_id = thread_config.get("collection_id")
+    user_id = thread_config.get("user_id")
     
     # Get files in the collection
-
-    collection_files = get_collection_files(collection_id, user_id) if collection_id and user_id else "No files available" #get_current_files()# 
-    logger.info(f"fetching collection for userid{user_id} and collection_id {collection_id}|| Results {collection_files}")
+    collection_files = get_collection_files(collection_id, user_id) if collection_id and user_id else "No files available"
+    logger.info(f"Fetching collection for userid {user_id} and collection_id {collection_id} || Results: {collection_files}")
+    
+    # Format using the prompt template
     return prompt.invoke({
         "collection_files": collection_files,
-        "messages": state["messages"]
+        "messages": state.get("messages", [])  # Safely get messages with default empty list
     })
 
 async def clean_tool_input(tool_input: str):