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 EventSourceResponse14from datetime import datetime15from fastapi import APIRouter16from langchain_core.runnables import RunnableConfig17from langchain_core.prompts import ChatPromptTemplate18from typing import Any19#from prompts import PRESENTATION_SYSTEM_PROMPT20 21router = APIRouter(22    prefix="/presentation",23    tags=["presentation"]24)25 26import json27 28@tool(response_format="content_and_artifact")29def plan(slides_json: str) -> tuple[str, dict]:30    """Create a presentation plan from a JSON string of slides (keys=slide numbers, values=content)."""31    try:32        slides = json.loads(slides_json)33        print(slides)34        return (35            f"Plan created with {len(slides)} slides: {', '.join(slides.keys())}.",36            {"slides_plan_json": slides_json}37        )38    except Exception as e:39        return (40            f"Invalid JSON format. Please provide a valid JSON string {str(e)[:100]}.",41            None42        )43 44@tool(response_format="content_and_artifact") 45def create_slide(slide_number: int, content: str, config: RunnableConfig) -> tuple[str, dict]:46    """Tool to create a full slide with the proposed content plan. Create content according to the plan, include images <!-- Title Slide -->47    <style scoped>48    </style>49    # Title50    <div class="subtitle"></div>51    <div class="logo"></div>52    <div class="date"></div>53    ![bg right:40%](https://picsum.photos/400/600)54    ---55    add styles to design a beutiful presentation, add a modern feel, add images using marp syntax and https://picsum.photos/200/300 (width & height)56    """57    # Integration with slide creation API or template would go here58    slide = {59        "number": slide_number,60        "content": content,61        "created_at": datetime.now().isoformat()62    }63    return (64        f"Slide {slide_number} created",65        {"slide": slide}66    )67 68@tool(parse_docstring=True)69def execute_python(expression: str) -> str:70    """Execute a python mathematic expression. Returns the result of the expression or an error message if execution fails.71 72    Args:73        expression: The python expression to execute.74    """75    try:76        result = eval(expression)77        return f"The result of the expression is {result}"78    except Exception as e:79        return f"Error executing the expression: {str(e)}"80 81memory = MemorySaver()82model = ChatOpenAI(model="gpt-4o-mini", streaming=True)83prompt = ChatPromptTemplate.from_messages([84    ("system", """You are a Presentation Creation Assistant. Your task is to help users create effective presentations.85    Follow these steps:86    1. First use the plan tool to create an outline of the presentation87    2. Wait for user input to proceed with the plan or update the current plan based on input.88    3. Then use create_slide tool for each slide in sequence89    4. Guide the user through the presentation creation process90    Today's date is {{datetime.now().strftime('%Y-%m-%d')}}"""),91    ("placeholder", "{messages}"),92])93 94def state_modifier(state) -> list[BaseMessage]:95    try:96        formatted_prompt = prompt.invoke({97            "messages": state["messages"]98        })99        print(state["messages"])100        return trim_messages(101            formatted_prompt,102            token_counter=len,103            max_tokens=16000,104            strategy="last",105            start_on="human",106            include_system=True,107            allow_partial=False,108        )109    except Exception as e:110        print(f"Error in state modifier: {str(e)}")111        return state["messages"]112 113# Create the agent with presentation tools114agent = create_react_agent(115    model,116    tools=[plan, create_slide, execute_python],117    checkpointer=memory,118    state_modifier=state_modifier,119)120 121class ChatInput(BaseModel):122    message: str123    thread_id: Optional[str] = None124 125@router.post("/chat")126async def chat(input_data: ChatInput):127    thread_id = input_data.thread_id or str(uuid.uuid4())128    129    config = {130        "configurable": {131            "thread_id": thread_id132        }133    }134    135    input_message = HumanMessage(content=input_data.message)136    137    async def generate():138        async for event in agent.astream_events(139            {"messages": [input_message]}, 140            config,141            version="v2"142        ):143            kind = event["event"]144            145            if kind == "on_chat_model_stream":146                content = event["data"]["chunk"].content147                if content:148                    yield f"{json.dumps({'type': 'token', 'content': content})}\n"149 150            elif kind == "on_tool_start":151                tool_input = str(event['data'].get('input', ''))152                yield f"{json.dumps({'type': 'tool_start', 'tool': event['name'], 'input': tool_input})}\n"153            154            elif kind == "on_tool_end":155                print(event['data'])156                tool_output = event['data'].get('output', '')157                artifact_output = tool_output.artifact if tool_output.artifact else None158                yield f"{json.dumps({'type': 'tool_end', 'tool': event['name'], 'output': tool_output.pretty_repr(), 'artifacts_data': artifact_output})}\n"159                print(tool_output.pretty_repr())160    return EventSourceResponse(161        generate(),162        media_type="text/event-stream"163    )164 165@router.get("/health")166async def health_check():167    return {"status": "healthy"}168 169 170app = FastAPI()171app.include_router(router)172 173if __name__ == "__main__":174    uvicorn.run(app, host="0.0.0.0", port=8000)175