Spaces:
Running
Running
| from __future__ import annotations | |
| from app.memory.history_trimmer import maybe_trim_history | |
| import json | |
| from typing import Any, Optional | |
| from langchain.agents import create_agent | |
| from langchain_nvidia_ai_endpoints import ChatNVIDIA | |
| from langgraph.checkpoint.redis.aio import AsyncRedisSaver | |
| from app.agent.prompts import SYSTEM_PROMPT | |
| from app.agent.state import AgentState | |
| from app.config import get_settings | |
| from app.tools.graph_tool import graph_query_tool | |
| from app.tools.sql_tool import sql_query_tool | |
| from app.tools.article_tool import get_article_detail | |
| from app.tools.hybrid_tool import hybrid_search_tool | |
| DEFAULT_MODEL = "z-ai/glm-5.2" | |
| _agent = None | |
| def build_agent(checkpointer: AsyncRedisSaver): | |
| """Called once at startup after the checkpointer is ready.""" | |
| global _agent | |
| settings = get_settings() | |
| llm = ChatNVIDIA( | |
| api_key=settings.nvidia_api_key, | |
| model=DEFAULT_MODEL, | |
| temperature=0, | |
| ) | |
| _agent = create_agent( | |
| model=llm, | |
| tools=[sql_query_tool, graph_query_tool, hybrid_search_tool], | |
| state_schema=AgentState, | |
| system_prompt=SYSTEM_PROMPT, | |
| checkpointer=checkpointer, | |
| ) | |
| return _agent | |
| def get_agent(): | |
| if _agent is None: | |
| raise RuntimeError( | |
| "Agent not initialized — call build_agent() at startup first") | |
| return _agent | |
| def _extract_sources(messages: list) -> list[dict]: | |
| sources: list[dict] = [] | |
| for m in messages: | |
| if getattr(m, "type", None) != "tool": | |
| continue | |
| try: | |
| payload = json.loads(m.content) if isinstance( | |
| m.content, str) else m.content | |
| except (json.JSONDecodeError, TypeError): | |
| continue | |
| refs = payload.get("source_refs") if isinstance( | |
| payload, dict) else None | |
| if refs: | |
| sources.extend(refs) | |
| seen: set[tuple] = set() | |
| deduped: list[dict] = [] | |
| for s in sources: | |
| key = (s.get("type"), s.get("id")) | |
| if key not in seen: | |
| seen.add(key) | |
| deduped.append(s) | |
| return deduped | |
| async def _enrich_sources(raw_refs: list[dict]) -> list[dict]: | |
| enriched: list[dict] = [] | |
| for ref in raw_refs: | |
| if ref.get("type") == "article" and ref.get("id") is not None: | |
| detail = await get_article_detail(int(ref["id"])) | |
| rows = detail.get("rows") or [] | |
| if rows: | |
| article = rows[0] | |
| enriched.append({ | |
| "type": "article", | |
| "id": article["id"], | |
| "title": article.get("title"), | |
| "url": article.get("url"), | |
| }) | |
| continue | |
| enriched.append(ref) | |
| return enriched | |
| async def run_agent(message: str, session_id: str, user_id: Optional[str]) -> dict[str, Any]: | |
| agent = get_agent() | |
| config = {"configurable": {"thread_id": session_id}} | |
| result = await agent.ainvoke( | |
| { | |
| "messages": [{"role": "user", "content": message}], | |
| "sources": [], | |
| "session_id": session_id, | |
| "user_id": user_id, | |
| }, | |
| config=config, | |
| ) | |
| final_message = result["messages"][-1] | |
| answer = final_message.content if isinstance( | |
| final_message.content, str) else str(final_message.content) | |
| raw_sources = _extract_sources(result["messages"]) | |
| sources = await _enrich_sources(raw_sources) | |
| await maybe_trim_history(agent, config) | |
| return {"answer": answer, "sources": sources} | |