Elevatics/code-chat-api

1import uuid2from fastapi import FastAPI3from fastapi.responses import StreamingResponse4from fastapi.middleware.cors import CORSMiddleware5from langchain_core.messages import BaseMessage, HumanMessage, trim_messages6from langchain_core.tools import tool7from langchain_openai import ChatOpenAI8from langgraph.checkpoint.memory import MemorySaver9from langgraph.prebuilt import create_react_agent10from pydantic import BaseModel11from typing import Optional12import json13from sse_starlette.sse import EventSourceResponse14import io15import sys16from contextlib import redirect_stdout, redirect_stderr17from langchain_core.runnables import RunnableConfig18import requests19import uvicorn20import re21from fastapi.staticfiles import StaticFiles22from langchain_core.runnables import RunnableConfig23from langchain_core.prompts import ChatPromptTemplate24from datetime import datetime25from presentation_api import router as presentation_router26from yf_docs import yf_docs27 28app = FastAPI()29app.include_router(presentation_router)30 31app.add_middleware(32    CORSMiddleware,33    allow_origins=["*"],34    allow_credentials=True,35    allow_methods=["*"],36    allow_headers=["*"],37)38 39app.mount("/chatui", StaticFiles(directory="static/chatui", html=True), name="index")40 41class CodeExecutionResult:42    def __init__(self, output: str, error: str = None):43        self.output = output44        self.error = error45 46API_URL = "https://vps-91587096.vps.ovh.us/api1"47 48@tool(response_format="content_and_artifact")49def execute_python(code: str, config: RunnableConfig):50    """Execute Python code in an jupyter notebook and return the output. The returned artifacts (if present) are automatically rendered in the UI and visible to the user. Available Libraries: plotly (default charting library),pandas,yfinance,numpy,geopandas,folium51    Args:52        code: Valid Python code with correct indentation and syntax including necessary imports.53    """54 55    thread_config = config.get("configurable", {})56    session_token = thread_config.get("thread_id", "test")57 58    headers = {59        'accept': 'application/json',60        'Content-Type': 'application/json'61    }62    data = {63        "session_token": session_token,64        "code": code65    }66    try:67        response = requests.post(68            f'{API_URL}/v0/execute',69            headers=headers,70            json=data71        )72        73        if response.status_code != 200:74            return (75                f"Error: Request failed with status code {response.status_code}. Response: {response.text}",76                None77            )78        79        # Get the response JSON80        response_json = response.json()81 82        # extract artifacts if they exist 83        artifacts_data = response_json.get("artifacts_data", {})84 85        # Create a clean response without artifacts86        execution_response = {87            "status": response_json.get("status"),88            "text": response_json.get("text"),89            "error_message": response_json.get("error_message"),90            "artifacts": response_json.get("artifacts")91        }92        93        return (94            f"Execution completed successfully: {json.dumps(execution_response)}",95            {"artifacts_data": artifacts_data} if artifacts_data else None96        )97        98    except Exception as e:99        return (f"Error executing code: {str(e)}", None)100 101save_yf_prompt = """102# Save downloaded data103tickers = yf.Tickers('MSFT AAPL')104data = tickers.history(period='6mo')105data.to_pickle('stock_data.pkl')106 107# Load it back later108data = pd.read_pickle('stock_data.pkl')109 110"""111 112memory = MemorySaver()113model = ChatOpenAI(model="gpt-4o", streaming=True)114prompt = ChatPromptTemplate.from_messages([115    ("system", f"You are a Data Visualization assistant.You have access to an jupyter notebook with access to internet for python code execution.\116    Your taks is to assist users with your data analysis and visualization expertise. Use only Plotly for creating visualizations and charts (Matplotlib is not available). Generated artifacts\117    are automatically rendered in the UI. Variables do not persist across tool calls, hence save any data to current directory you want to use in the next tool call to a file.(use file_name and do not add file path, as you have only permission to edit current folder) {save_yf_prompt} Today's date is \118    {datetime.now().strftime('%Y-%m-%d')}. Format you responses Beutifully using markdown tables, paragraphs, lists etc. Saved files are not accessible to user.The current folder contains the following files: {{collection_files}} {yf_docs}"),119    ("placeholder", "{messages}"),120])121 122def state_modifier(state) -> list[BaseMessage]:123    collection_files = "None"124    try:125        formatted_prompt = prompt.invoke({126            "collection_files": collection_files,127            "messages": state["messages"]128        })129        print(state["messages"])130        return trim_messages(131            formatted_prompt,132            token_counter=len,133            max_tokens=16000,134            strategy="last",135            start_on="human",136            include_system=True,137            allow_partial=False,138        )139        140    except Exception as e:141        print(f"Error in state modifier: {str(e)}")142        return state["messages"]143 144# Create the agent with the Python execution tool145agent = create_react_agent(146    model,147    tools=[execute_python],148    checkpointer=memory,149    state_modifier=state_modifier,150)151 152class ChatInput(BaseModel):153    message: str154    thread_id: Optional[str] = None155 156@app.post("/chat")157async def chat(input_data: ChatInput):158    thread_id = input_data.thread_id or str(uuid.uuid4())159    160    config = {161        "configurable": {162            "thread_id": thread_id163        }164    }165    166    input_message = HumanMessage(content=input_data.message)167    168    async def generate():169        async for event in agent.astream_events(170            {"messages": [input_message]}, 171            config,172            version="v2"173        ):174            kind = event["event"]175            176            if kind == "on_chat_model_stream":177                content = event["data"]["chunk"].content178                if content:179                    yield f"{json.dumps({'type': 'token', 'content': content})}\n"180 181            elif kind == "on_tool_start":182                tool_input = event['data'].get('input', '')183                yield f"{json.dumps({'type': 'tool_start', 'tool': event['name'], 'input': tool_input})}\n"184            185            elif kind == "on_tool_end":186                tool_output = event['data'].get('output', '').content187                artifact_output = event['data'].get('output', '').artifact.get('artifacts_data') if event['data'].get('output', '').artifact else None188                yield f"{json.dumps({'type': 'tool_end', 'tool': event['name'], 'output': tool_output, 'artifacts_data': artifact_output})}\n"189                #print(f"{json.dumps({'type': 'tool_end', 'tool': event['name'], 'output': tool_output, 'artifacts_data': artifact_output})}\n")190    return EventSourceResponse(191        generate(),192        media_type="text/event-stream"193    )194 195@app.get("/health")196async def health_check():197    return {"status": "healthy"}198 199if __name__ == "__main__":200    uvicorn.run(app, host="0.0.0.0", port=7860)