Create chat_langraph.py

src/chat_langraph.py

@@ -0,0 +1,126 @@
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_core.messages import BaseMessage, ToolMessage, AIMessage, SystemMessage, HumanMessage
+from langgraph.graph import StateGraph, add_messages, START, END
+from langgraph.checkpoint.sqlite import SqliteSaver
+from typing import TypedDict, Annotated, List
+from langchain_core.tools import tool
+from langgraph.prebuilt.tool_node import ToolNode
+import sqlite3
+import subprocess
+import requests
+from datetime import datetime
+
+class chatstate(TypedDict):
+    messages: Annotated[List[BaseMessage], add_messages]
+
+
+api = "AIzaSyA5zvErF4vUmAoslVzkOBUfvSCSoW0vjEA"
+LANGSEARCH_API_KEY = "sk-f1a8f996f9e44b43adf9943e43e8582b"
+
+llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0.2, api_key=api)
+
+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 . 
+Rules:
+1. Prefer text answers and code when examples/explanations are asked.
+2. Explicit requests to create/run files → call appropriate tool.
+3. Avoid destructive commands without confirmation.
+4. Keep tool inputs minimal.
+Tone: concise, helpful, decisive.
+"""
+)
+
+conn = sqlite3.connect("chatbot.db", check_same_thread=False)
+checkpointer = SqliteSaver(conn=conn)
+
+
+@tool
+def add(a: int, b: int):
+    return a + b
+
+
+@tool
+def reverse(string: str):
+    return string[::-1]
+
+
+@tool
+def evaluate(string: str):
+    return eval(string)
+
+
+@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}"
+
+
+@tool
+def run_cmd_command(command: str) -> str:
+    """Run a safe shell command on linux """
+    try:
+        result = subprocess.run(command, shell=True, check=True, text=True, capture_output=True)
+        return result.stdout
+    except subprocess.CalledProcessError as e:
+        return f"Error: {e}"
+
+
+@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):
+    return "end" if state["messages"][-1].content == "end" else "llmresponse"
+
+
+def input_node(state: chatstate):
+    return {"messages": state["messages"]}
+
+
+def llmresponse(state: chatstate):
+    response = llm.invoke(state["messages"])
+    return {"messages": [response]}
+
+
+def checktool(state: chatstate):
+    last_msg = state["messages"][-1]
+    if hasattr(last_msg, "tool_calls") and last_msg.tool_calls:
+        return "tool_node"
+    return "end"
+
+
+tools = [add, reverse, evaluate, run_cmd_command, search_tool, write_file]
+tool_node = ToolNode(tools=tools)
+llm = llm.bind_tools(tools)
+
+graph = StateGraph(chatstate)
+graph.add_node("input_node", input_node)
+graph.add_node("llmresponse", llmresponse)
+graph.add_node("tool_node", tool_node)
+graph.add_edge(START, "input_node")
+graph.add_edge("input_node", "llmresponse")
+graph.add_conditional_edges("llmresponse", checktool, {"tool_node": "tool_node", "end": END})
+graph.add_edge("tool_node", "llmresponse")
+
+workflow = graph.compile(checkpointer=checkpointer)
+
+
+def get_all_chat_ids():
+    s = set()
+    for chkpoint in checkpointer.list(None):
+        s.add(chkpoint.config.get("configurable").get("thread_id"))
+    return list(s)