Create main.py

main.py

@@ -0,0 +1,186 @@
+import uuid
+from fastapi import FastAPI
+from fastapi.responses import StreamingResponse
+from fastapi.middleware.cors import CORSMiddleware
+from langchain_core.messages import BaseMessage, HumanMessage, trim_messages
+from langchain_core.tools import tool
+from langchain_openai import ChatOpenAI
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.prebuilt import create_react_agent
+from pydantic import BaseModel
+from typing import Optional
+import json
+from sse_starlette.sse import EventSourceResponse
+import io
+import sys
+from contextlib import redirect_stdout, redirect_stderr
+from langchain_core.runnables import RunnableConfig
+import requests
+import uvicorn
+import re
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+class CodeExecutionResult:
+    def __init__(self, output: str, error: str = None):
+        self.output = output
+        self.error = error
+
+API_URL = "https://pvanand-code-execution-files-v4.hf.space"
+
+@tool
+def execute_python(code: str) -> str:
+    """Execute Python code in an IPython interactiveshell and return the output.
+    
+    Args:
+        code: The Python code to execute
+
+    Available Libraries:
+    # Use plotly as the default charting library
+    # While using yfinance to pull stock data, Always clean the multiindex columns as this might cause issues in plotting plotly charts
+    # Remove the ticker level from columns if it exists
+    yf_data = yf.download(symbol, start=start_date, end=end_date)
+    if isinstance(yf_data.columns, pd.MultiIndex):
+        yf_data.columns = yf_data.columns.get_level_values(0)
+
+    matplotlib
+    pandas
+    plotly
+    groq
+    yfinance
+    numpy
+    seaborn
+    numpy
+    scikit-learn
+    statsmodels
+    geopandas
+    folium
+    fpdf
+    kaleido
+    scipy
+    geopy
+    mapbox
+
+    Artifacts are automatically rendered in the UI hence no need to provide links to them.
+    """
+
+    #print(config)
+
+    headers = {
+        'accept': 'application/json',
+        'Content-Type': 'application/json'
+    }
+    data = {
+        "session_token": "test12345", #config.configurable.get("thread_id", "test"),
+        "code": code
+    }
+    response = requests.post(
+        f'{API_URL}/v0/execute',
+        headers=headers,
+        data=json.dumps(data)
+    )
+
+    if response.status_code != 200:
+        return f"Error: Request failed with status code {response.status_code}. Response: {response.text}"
+    else:
+        response_json = response.json()
+        return f"data: {json.dumps(response_json)} \ndata:"
+
+# Configure the memory and model"
+memory = MemorySaver()
+model = ChatOpenAI(model="gpt-4o-mini", streaming=True)
+
+def state_modifier(state) -> list[BaseMessage]:
+    return trim_messages(
+        state["messages"],
+        token_counter=len,
+        max_tokens=16000,
+        strategy="last",
+        start_on="human",
+        include_system=True,
+        allow_partial=False,
+    )
+
+# Create the agent with the Python execution tool
+agent = create_react_agent(
+    model,
+    tools=[execute_python],
+    checkpointer=memory,
+    state_modifier=state_modifier,
+)
+
+class ChatInput(BaseModel):
+    message: str
+    thread_id: Optional[str] = None
+
+@app.post("/chat")
+async def chat(input_data: ChatInput):
+    thread_id = input_data.thread_id or str(uuid.uuid4())
+    
+    config = {
+        "configurable": {
+            "thread_id": thread_id
+        }
+    }
+    
+    input_message = HumanMessage(content=input_data.message)
+    
+    async def generate():
+        async for event in agent.astream_events(
+            {"messages": [input_message]}, 
+            config,
+            version="v2"
+        ):
+            kind = event["event"]
+            
+            if kind == "on_chat_model_stream":
+                content = event["data"]["chunk"].content
+                if content:
+                    yield f"{json.dumps({'type': 'token', 'content': content})}\n"
+
+            elif kind == "on_tool_start":
+                tool_input = event['data'].get('input', '')
+                yield f"{json.dumps({'type': 'tool_start', 'tool': event['name'], 'input': tool_input})}\n"
+            
+            elif kind == "on_tool_end":
+                tool_output = event['data'].get('output', '').content
+                #print(type(tool_output))
+                #print(dir(tool_output))
+                #print the keys
+                pattern = r'data: (.*?)\ndata:'
+                match = re.search(pattern, tool_output)
+                print(tool_output)
+
+                if match:
+                    tool_output_json = match.group(1).strip()
+                    try:
+                        tool_output = json.loads(tool_output_json)
+                        if "artifacts" in tool_output:
+                            for artifact in tool_output["artifacts"]:
+                                artifact_content = requests.get(f"{API_URL}/artifact/{artifact['artifact_id']}").content
+                                print(artifact_content)
+                                tool_output["artifacts"][artifact["artifact_id"]] = artifact_content
+                    except Exception as e:
+                        print(e)
+                        print("Error parsing tool output as json: ", tool_output)
+                else:
+                    print("No match found in tool output")  
+                yield f"{json.dumps({'type': 'tool_end', 'tool': event['name'], 'output': tool_output})}\n"
+    return EventSourceResponse(
+        generate(),
+        media_type="text/event-stream"
+    )
+
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=9000)