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 for creating visualizations

    plotly
    pandas
    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. Use Plotly for creating visualizations. 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"
    try:
        formatted_prompt = prompt.invoke({
            "collection_files": collection_files,
            "messages": state["messages"]
        })
        
        return trim_messages(
            formatted_prompt,
            token_counter=len,
            max_tokens=16000,
            strategy="last",
            start_on="human",
            include_system=True,
            allow_partial=False,
        )
    print(state["messages"])
    except Exception as e:
        print(f"Error in state modifier: {str(e)}")
        return state["messages"]

# 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"

            if kind == "on_tool_start":
                print(f"Debug - Tool being called: {event['name']}")  # Add this debug line
                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)