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
from fastapi.staticfiles import StaticFiles
from langchain_core.runnables import RunnableConfig
from langchain_core.prompts import ChatPromptTemplate
from datetime import datetime

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.mount("/chatui", StaticFiles(directory="static/chatui", html=True), name="index")

class CodeExecutionResult:
    def __init__(self, output: str, error: str = None):
        self.output = output
        self.error = error

API_URL = "https://pvanand-code-execution-files-v5.hf.space"

@tool
def execute_python(code: str, config: RunnableConfig):
    """Execute Python code in an IPython interactiveshell and return the output.
       The returned artifacts (if present) are automatically rendered in the UI and visible to the user. 
    Args:
        code: Valid Python code with correct indentation and syntax including necessary imports.

    Available Libraries:
    # Use plotly as the default charting library

    matplotlib
    pandas
    plotly
    groq
    yfinance
    numpy
    seaborn
    numpy
    scikit-learn
    statsmodels
    geopandas
    folium
    fpdf
    kaleido
    scipy
    geopy
    mapbox
    """

    #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)
prompt = ChatPromptTemplate.from_messages([
    ("system", f"You are a Data Visualization assistant.You have access to a jupyter client with access to internet for python code execution. Your taks is to assist users with your data analysis and visualization expertise. Today's date is {datetime.now().strftime('%Y-%m-%d')}. The current folder contains the following files: {{collection_files}}"),
    ("placeholder", "{messages}"),
])

def state_modifier(state) -> list[BaseMessage]:
    collection_files = "None"
    # Format the prompt with the current state
    formatted_prompt = prompt.invoke({
        "collection_files": collection_files,
        "messages": state["messages"]
    })
    
    # Trim the messages
    return trim_messages(
        formatted_prompt,
        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=7860)