Spaces:
Running
Running
File size: 4,031 Bytes
0b54eef |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 |
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
from datetime import datetime
from fastapi import APIRouter
router = APIRouter(
prefix="/presentation",
tags=["presentation"]
)
@tool
def plan(input: dict) -> str:
"""Create a presentation plan with numbered slides and their descriptions.
Args:
input: Dictionary containing presentation details
Returns:
A dictionary with slide numbers as keys and descriptions as values
"""
return "plan created"
@tool
def create_slide(slideno: int, content: str) -> str:
"""Create a single presentation slide.
Args:
slideno: The slide number to create
content: The content for the slide
Returns:
Confirmation of slide creation
"""
return f"slide {slideno} created"
memory = MemorySaver()
model = ChatOpenAI(model="gpt-4-turbo-preview", streaming=True)
prompt = ChatPromptTemplate.from_messages([
("system", """You are a Presentation Creation Assistant. Your task is to help users create effective presentations.
Follow these steps:
1. First use the plan tool to create an outline of the presentation
2. Then use create_slide tool for each slide in sequence
3. Guide the user through the presentation creation process
Today's date is {datetime.now().strftime('%Y-%m-%d')}"""),
("placeholder", "{messages}"),
])
def state_modifier(state) -> list[BaseMessage]:
try:
formatted_prompt = prompt.invoke({
"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,
)
except Exception as e:
print(f"Error in state modifier: {str(e)}")
return state["messages"]
# Create the agent with presentation tools
agent = create_react_agent(
model,
tools=[plan, create_slide],
checkpointer=memory,
state_modifier=state_modifier,
)
class ChatInput(BaseModel):
message: str
thread_id: Optional[str] = None
@router.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', '')
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) |