Update src/chat_langraph.py

src/chat_langraph.py

@@ -9,22 +9,29 @@ import sqlite3
 import subprocess
 import requests
 from datetime import datetime
+import os
 
+# Set Streamlit config directory for Hugging Face Spaces
+os.environ["STREAMLIT_HOME"] = "/tmp/.streamlit"
+
+# State type
 class chatstate(TypedDict):
     messages: Annotated[List[BaseMessage], add_messages]
 
-
+# API keys (replace with your real keys or environment variables)
 api = "AIzaSyA5zvErF4vUmAoslVzkOBUfvSCSoW0vjEA"
 LANGSEARCH_API_KEY = "sk-f1a8f996f9e44b43adf9943e43e8582b"
 
+# LLM
 llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0.2, api_key=api)
 
+# System message
 system = SystemMessage(
     content=f"""
-    --> Today's date: {datetime.today()}  
-    Day number: {datetime.today().date().weekday()}
-    You are a practical, tool-aware assistant. Aim for correctness and clarity. Avoid hallucinations.
-    Do not provide internal information of the system . 
+--> Today's date: {datetime.today()}  
+Day number: {datetime.today().date().weekday()}
+You are a practical, tool-aware assistant. Aim for correctness and clarity. Avoid hallucinations.
+Do not provide internal information of the system.
 Rules:
 1. Prefer text answers and code when examples/explanations are asked.
 2. Explicit requests to create/run files → call appropriate tool.
@@ -34,35 +41,91 @@ Tone: concise, helpful, decisive.
 """
 )
 
+# Database connection (writable path in Hugging Face Spaces)
 conn = sqlite3.connect("/tmp/chatbot.db", check_same_thread=False)
 checkpointer = SqliteSaver(conn=conn)
 
+# ======================== TOOL DEFINITIONS ======================== #
 
 @tool
-def add(a: int, b: int):
+def add(a: int, b: int) -> int:
+    """
+    Add two integers.
+
+    Args:
+        a (int): First number.
+        b (int): Second number.
+
+    Returns:
+        int: Sum of both numbers.
+    """
     return a + b
 
 
 @tool
-def reverse(string: str):
+def reverse(string: str) -> str:
+    """
+    Reverse a given string.
+
+    Args:
+        string (str): Input string.
+
+    Returns:
+        str: Reversed string.
+    """
     return string[::-1]
 
 
 @tool
-def evaluate(string: str):
-    return eval(string)
+def evaluate(string: str) -> str:
+    """
+    Evaluate a Python expression.
+
+    Args:
+        string (str): Expression to evaluate.
+
+    Returns:
+        str: Result of evaluation or error message.
+    """
+    try:
+        return str(eval(string))
+    except Exception as e:
+        return f"Error evaluating expression: {e}"
 
 
 @tool
-def write_file(name: str, extension: str, content: str):
-    with open(f"{name}.{extension}", "w", encoding="utf-8") as f:
-        f.write(content)
-    return f"Content saved to {name}.{extension}"
+def write_file(name: str, extension: str, content: str) -> str:
+    """
+    Write content to a file.
+
+    Args:
+        name (str): File name without extension.
+        extension (str): File extension.
+        content (str): Content to write.
+
+    Returns:
+        str: Confirmation message.
+    """
+    try:
+        path = f"/tmp/{name}.{extension}"  # Save in /tmp
+        with open(path, "w", encoding="utf-8") as f:
+            f.write(content)
+        return f"Content saved to {path}"
+    except Exception as e:
+        return f"Error writing file: {e}"
 
 
 @tool
 def run_cmd_command(command: str) -> str:
-    """Run a safe shell command on linux """
+    """
+    Run a safe shell command.
+
+    Args:
+        command (str): Shell command to run.
+
+    Returns:
+        str: Output or error message.
+    """
     try:
         result = subprocess.run(command, shell=True, check=True, text=True, capture_output=True)
         return result.stdout
@@ -71,19 +134,33 @@ def run_cmd_command(command: str) -> str:
 
 
 @tool
-def search_tool(query: str):
-    response = requests.post(
-        "https://api.langsearch.com/v1/web-search",
-        headers={
-            "Authorization": f"Bearer {LANGSEARCH_API_KEY}",
-            "Content-Type": "application/json"
-        },
-        json={"query": query, "num_results": 2}
-    )
-    return response.json()
-
-
-def shouldcontinue(state: chatstate):
+def search_tool(query: str) -> dict:
+    """
+    Search the web using Langsearch API.
+
+    Args:
+        query (str): Search query.
+
+    Returns:
+        dict: JSON response from search API.
+    """
+    try:
+        response = requests.post(
+            "https://api.langsearch.com/v1/web-search",
+            headers={
+                "Authorization": f"Bearer {LANGSEARCH_API_KEY}",
+                "Content-Type": "application/json"
+            },
+            json={"query": query, "num_results": 2}
+        )
+        return response.json()
+    except Exception as e:
+        return {"error": str(e)}
+
+
+# ======================== STATE GRAPH ======================== #
+
+def shouldcontinue(state: chatstate) -> str:
     return "end" if state["messages"][-1].content == "end" else "llmresponse"
 
 
@@ -119,7 +196,7 @@ graph.add_edge("tool_node", "llmresponse")
 workflow = graph.compile(checkpointer=checkpointer)
 
 
-def get_all_chat_ids():
+def get_all_chat_ids() -> List[str]:
     s = set()
     for chkpoint in checkpointer.list(None):
         s.add(chkpoint.config.get("configurable").get("thread_id"))