64 lines
2.0 KiB
Python
64 lines
2.0 KiB
Python
"""LangGraph definition for Gumbo (load from MCP FS -> summarize via LiteLLM)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from typing import TypedDict
|
|
|
|
from langchain_core.messages import HumanMessage, SystemMessage
|
|
from langchain_openai import ChatOpenAI
|
|
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
|
from langgraph.graph import END, StateGraph
|
|
|
|
from gumbo.mcp_fs import fetch_text_via_mcp
|
|
|
|
|
|
class GumboState(TypedDict):
|
|
object_key: str
|
|
source_text: str
|
|
summary: str
|
|
|
|
|
|
async def _load(state: GumboState) -> dict:
|
|
text = await fetch_text_via_mcp(state["object_key"])
|
|
return {"source_text": text}
|
|
|
|
|
|
async def _summarize(state: GumboState) -> dict:
|
|
llm = ChatOpenAI(
|
|
model=os.environ.get("GUMBO_LLM_MODEL", "ollama-qwen"),
|
|
api_key=os.environ["LITELLM_API_KEY"],
|
|
base_url=os.environ["LITELLM_BASE_URL"],
|
|
temperature=0.2,
|
|
)
|
|
system = SystemMessage(
|
|
content="You are Gumbo, a precise documentation summarizer. Return a concise markdown summary."
|
|
)
|
|
user = HumanMessage(
|
|
content=f"Summarize the following document:\n\n{state['source_text']}",
|
|
)
|
|
resp = await llm.ainvoke([system, user])
|
|
summary = resp.content if isinstance(resp.content, str) else str(resp.content)
|
|
return {"summary": summary}
|
|
|
|
|
|
def build_graph_builder() -> StateGraph:
|
|
builder = StateGraph(GumboState)
|
|
builder.add_node("load", _load)
|
|
builder.add_node("summarize", _summarize)
|
|
builder.set_entry_point("load")
|
|
builder.add_edge("load", "summarize")
|
|
builder.add_edge("summarize", END)
|
|
return builder
|
|
|
|
|
|
async def run_gumbo(object_key: str, thread_id: str, conn_string: str) -> GumboState:
|
|
builder = build_graph_builder()
|
|
async with AsyncPostgresSaver.from_conn_string(conn_string) as checkpointer:
|
|
await checkpointer.setup()
|
|
graph = builder.compile(checkpointer=checkpointer)
|
|
return await graph.ainvoke(
|
|
{"object_key": object_key, "source_text": "", "summary": ""},
|
|
config={"configurable": {"thread_id": thread_id}},
|
|
)
|