main.py

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)