Update main.py

main.py

@@ -178,10 +178,19 @@ prompt = ChatPromptTemplate.from_messages([
     ("placeholder", "{messages}"),
 ])
 
+import requests
+from requests.exceptions import RequestException, Timeout
+import logging
+from typing import Optional
+
+
 def get_collection_files(collection_id: str, user_id: str) -> str:
-    """Get list of files in the specified collection using the external API"""
+    """
+    Synchronously get list of files in the specified collection using the external API
+    with proper timeout and error handling.
+    """
     try:
-        url = f"https://pvanand-documind-api-v2.hf.space/rag/get_collection_files"
+        url = "https://pvanand-documind-api-v2.hf.space/rag/get_collection_files"
         params = {
             "collection_id": collection_id,
             "user_id": user_id
@@ -189,16 +198,28 @@ def get_collection_files(collection_id: str, user_id: str) -> str:
         headers = {
             'accept': 'application/json'
         }
-        response = requests.post(url, params=params, headers=headers, data='')
+        
+        logger.debug(f"Requesting collection files for user {user_id}, collection {collection_id}")
+        
+        # Set timeout to 5 seconds
+        response = requests.post(url, params=params, headers=headers, data='', timeout=5)
         
         if response.status_code == 200:
+            logger.info(f"Successfully retrieved collection files: {response.text[:100]}...")
             return response.text
         else:
-            logger.error(f"Error from API: {response.text}")
-            return f"Error getting files: {response.text}"
+            logger.error(f"API error (status {response.status_code}): {response.text}")
+            return f"Error fetching files (status {response.status_code})"
+            
+    except Timeout:
+        logger.error("Timeout while fetching collection files")
+        return "Error: Request timed out"
+    except RequestException as e:
+        logger.error(f"Network error fetching collection files: {str(e)}")
+        return f"Error: Network issue - {str(e)}"
     except Exception as e:
-        logger.error(f"Error getting collection files: {str(e)}")
-        return f"Error getting files: {str(e)}"
+        logger.error(f"Error fetching collection files: {str(e)}", exc_info=True)
+        return f"Error fetching files: {str(e)}"
 
 def format_for_model(state: dict, config: Optional[RunnableConfig] = None) -> list[BaseMessage]:
     """
@@ -216,15 +237,28 @@ def format_for_model(state: dict, config: Optional[RunnableConfig] = None) -> li
     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"
-    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.get("messages", [])  # Safely get messages with default empty list
-    })
+    try:
+        # Get files in the collection with timeout protection
+        if collection_id and user_id:
+            collection_files = get_collection_files(collection_id, user_id)
+        else:
+            collection_files = "No files available"
+            
+        logger.info(f"Fetching collection for userid {user_id} and collection_id {collection_id} || Results: {collection_files[:100]}...")
+        
+        # Format using the prompt template
+        return prompt.invoke({
+            "collection_files": collection_files,
+            "messages": state.get("messages", [])
+        })
+        
+    except Exception as e:
+        logger.error(f"Error in format_for_model: {str(e)}", exc_info=True)
+        # Return a basic format if there's an error
+        return prompt.invoke({
+            "collection_files": "Error fetching files",
+            "messages": state.get("messages", [])
+        })
 
 async def clean_tool_input(tool_input: str):
     # Use regex to parse the first key and value