Engineering AI-First SaaS: Architecture Patterns for Autonomous Product Features

import asyncio
import json
import uuid
from typing import AsyncGenerator

from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from redis.asyncio import Redis
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessageChunk
from langgraph.graph import StateGraph, MessagesState, START, END
from langgraph.prebuilt import ToolNode
from langchain_core.tools import tool

app = FastAPI()
redis = Redis(host="localhost", port=6379, decode_responses=True)


# --- Tool definitions ---
@tool
async def search_knowledge_base(query: str) -> str:
    """Search the enterprise knowledge base for relevant information."""
    # Stub: in production this calls your hybrid retrieval pipeline
    await asyncio.sleep(0.5)  # simulate retrieval latency
    return f"Found 3 relevant documents for: {query}"


@tool
async def draft_summary(content: str) -> str:
    """Draft a structured summary from retrieved content."""
    await asyncio.sleep(0.3)
    return f"Summary: {content[:200]}..."


# --- LangGraph agent setup ---
tools = [search_knowledge_base, draft_summary]
tool_node = ToolNode(tools)
llm = ChatOpenAI(model="gpt-4o", streaming=True).bind_tools(tools)


def should_continue(state: MessagesState):
    last = state["messages"][-1]
    return "tools" if last.tool_calls else END


async def call_llm(state: MessagesState):
    response = await llm.ainvoke(state["messages"])
    return {"messages": [response]}


graph = StateGraph(MessagesState)
graph.add_node("agent", call_llm)
graph.add_node("tools", tool_node)
graph.add_edge(START, "agent")
graph.add_conditional_edges("agent", should_continue)
graph.add_edge("tools", "agent")
workflow = graph.compile()


# --- Request / job models ---
class AnalysisRequest(BaseModel):
    query: str
    user_id: str


class JobResponse(BaseModel):
    job_id: str
    stream_url: str


# --- Async worker that runs the agent and publishes events to Redis ---
async def run_agent_job(job_id: str, query: str, user_id: str) -> None:
    channel = f"job:{job_id}:events"
    try:
        async for event in workflow.astream_events(
            {"messages": [HumanMessage(content=query)]},
            version="v2",
        ):
            kind = event["event"]
            if kind == "on_chat_model_stream":
                chunk: AIMessageChunk = event["data"]["chunk"]
                if chunk.content:
                    await redis.publish(channel, json.dumps({"type": "token", "content": chunk.content}))
            elif kind == "on_tool_start":
                tool_name = event["name"]
                await redis.publish(channel, json.dumps({"type": "tool_start", "tool": tool_name}))
            elif kind == "on_tool_end":
                await redis.publish(channel, json.dumps({"type": "tool_end", "tool": event["name"]}))

        await redis.publish(channel, json.dumps({"type": "done"}))
    except Exception as e:
        await redis.publish(channel, json.dumps({"type": "error", "message": str(e)}))


# --- HTTP endpoints ---
@app.post("/analyze", response_model=JobResponse)
async def start_analysis(req: AnalysisRequest):
    job_id = str(uuid.uuid4())
    # Fire-and-forget: run agent in background
    asyncio.create_task(run_agent_job(job_id, req.query, req.user_id))
    return JobResponse(job_id=job_id, stream_url=f"/analyze/{job_id}/stream")


async def _sse_generator(job_id: str) -> AsyncGenerator[str, None]:
    channel = f"job:{job_id}:events"
    pubsub = redis.pubsub()
    await pubsub.subscribe(channel)
    try:
        async for message in pubsub.listen():
            if message["type"] != "message":
                continue
            data = json.loads(message["data"])
            yield f"data: {json.dumps(data)}\n\n"
            if data["type"] in ("done", "error"):
                break
    finally:
        await pubsub.unsubscribe(channel)


@app.get("/analyze/{job_id}/stream")
async def stream_analysis(job_id: str):
    return StreamingResponse(
        _sse_generator(job_id),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )