"""
Retrieval Augmented Generation (RAG) for conference abstracts.
This module provides RAG functionality using Pydantic AI as the agent framework.
The agent uses tools from the MCP server to search papers, analyze topics, and
understand trends, then generates contextual responses.
"""
from datetime import datetime
import json
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Dict, Optional, Any, Callable
import httpx2
from pydantic_ai import Agent, RunContext, Tool
from pydantic_ai.messages import ModelMessage, ModelRequest
from pydantic_ai.models.openai import OpenAIChatModel, OpenAIChatModelSettings
from pydantic_ai.providers.openai import OpenAIProvider
from abstracts_explorer.config import get_config
from abstracts_explorer.mcp_server import (
search_papers as mcp_search_papers,
get_conference_topics as mcp_get_conference_topics,
get_topic_evolution as mcp_get_topic_evolution,
analyze_topic_relevance as mcp_analyze_topic_relevance,
get_cluster_visualization as mcp_get_cluster_visualization,
get_paper_details as mcp_get_paper_details,
summarize_interesting_papers as mcp_summarize_interesting_papers,
recommend_related_papers as mcp_recommend_related_papers,
)
from abstracts_explorer.mcp_tools import format_tool_result_for_llm, _abbreviate_result
logger = logging.getLogger(__name__)
[docs]
class RAGError(Exception):
"""Exception raised for RAG-related errors."""
pass
[docs]
@dataclass
class RAGDeps:
"""
Dependencies for RAG agent tools.
Carries shared state for capturing tool results during agent runs,
allowing post-processing (e.g. paper extraction) after the agent completes.
Attributes
----------
tool_results : list of dict
Accumulated tool call results from the current agent run.
Each entry has 'name' (tool name) and 'raw_result' (JSON string).
conferences : list of str
Default conferences for MCP clustering tools (from web UI selection).
years : list of int
Default years for MCP clustering tools (from web UI selection).
paper_priorities : dict
Paper ratings from the web UI (``localStorage.paperPriorities``).
Keys are paper UIDs; values are either star ratings or dicts with
``priority`` and ``searchTerm`` fields.
"""
tool_results: List[Dict[str, Any]] = field(default_factory=list)
conferences: List[str] = field(default_factory=list)
years: List[int] = field(default_factory=list)
n_results: Optional[int] = None
paper_priorities: Dict[str, Any] = field(default_factory=dict)
# ---------------------------------------------------------------------------
# Tool wrapper functions
# ---------------------------------------------------------------------------
# These thin wrappers call the MCP server functions and capture the raw result
# in the deps for later paper extraction. They return a human-readable
# summary that goes back to the LLM as tool output.
# ---------------------------------------------------------------------------
def _run_rag_tool(
ctx: RunContext[RAGDeps],
tool_name: str,
mcp_func: Callable[..., str],
kwargs: Dict[str, Any],
) -> str:
"""
Execute an MCP tool and record its result for post-processing.
Parameters
----------
ctx : RunContext[RAGDeps]
Agent context with dependencies.
tool_name : str
Name of the tool as exposed to the LLM.
mcp_func : callable
MCP server function to invoke.
kwargs : dict
Keyword arguments for *mcp_func*.
Returns
-------
str
Human-readable tool result for the LLM.
"""
logger.info("Tool call: %s(%s)", tool_name, kwargs)
raw = mcp_func(**kwargs)
logger.info("Tool result: %s → %s", tool_name, _abbreviate_result(raw))
ctx.deps.tool_results.append({"name": tool_name, "raw_result": raw})
return format_tool_result_for_llm(tool_name, raw)
def _tool_search_papers(
ctx: RunContext[RAGDeps],
topic_keywords: str,
n_results: Optional[int] = None,
years: Optional[List[int]] = None,
conference: Optional[str] = None,
limit: Optional[int] = None,
) -> str:
"""Search for papers about a specific topic using semantic search.
Use this tool to find papers related to a research topic. Returns
matching papers with titles, abstracts, and relevance scores.
A conference must be specified. If not provided by the user, the
currently selected conference from the web UI is used automatically.
The currently selected year is also applied automatically when not
explicitly provided.
Parameters
----------
ctx : RunContext[RAGDeps]
Agent context with dependencies.
topic_keywords : str
Keywords describing the topic to search for.
n_results : int, optional
Number of papers to return (default: 10, or session default).
years : list of int, optional
Filter by publication years (e.g. [2024, 2025]).
When omitted, the default year(s) from the current session are used.
conference : str, optional
Filter by conference name (e.g. "NeurIPS", "ICLR").
When omitted, the default conference from the current session is used.
limit : int, optional
Alias for n_results.
"""
count = n_results if n_results is not None else (limit if limit is not None else (ctx.deps.n_results or 10))
kwargs: Dict[str, Any] = {"topic_keywords": topic_keywords, "n_results": count}
# Determine years: explicit param > deps default
if years is not None:
kwargs["years"] = years
elif ctx.deps.years:
kwargs["years"] = ctx.deps.years
# Determine conference: explicit param > deps default
if conference:
kwargs["conference"] = conference
elif ctx.deps.conferences:
kwargs["conference"] = ctx.deps.conferences[0]
return _run_rag_tool(ctx, "search_papers", mcp_search_papers, kwargs)
def _tool_get_conference_topics(
ctx: RunContext[RAGDeps],
conference: Optional[str] = None,
) -> str:
"""Get the main research topics of a conference.
Returns the key research topics covered at the conference, each with
a descriptive name, representative keywords, paper count, and example
paper titles.
A conference must be specified. If not provided by the user, the
currently selected conference from the web UI is used automatically.
Parameters
----------
ctx : RunContext[RAGDeps]
Agent context with dependencies.
conference : str, optional
Conference name (e.g. "NeurIPS").
When omitted, the default conference from the current session is used.
"""
# Determine conferences: explicit param > deps default
if conference:
conferences = [conference]
elif ctx.deps.conferences:
conferences = ctx.deps.conferences
else:
conferences = None
years = ctx.deps.years if ctx.deps.years else None
kwargs = {"conferences": conferences, "years": years}
return _run_rag_tool(ctx, "get_conference_topics", mcp_get_conference_topics, kwargs)
def _tool_get_topic_evolution(
ctx: RunContext[RAGDeps],
topic_keywords: str,
conferences: Optional[list[str]] = None,
start_year: Optional[int] = None,
end_year: Optional[int] = None,
) -> str:
"""Analyze how a topic has evolved over the years.
Use this tool to understand trends and year-over-year changes
in a research topic. Returns paper counts and relative percentages
by year for one or more conferences.
At least one conference must be specified. If not provided by the user,
the currently selected conferences from the web UI are used automatically.
Parameters
----------
ctx : RunContext[RAGDeps]
Agent context with dependencies.
topic_keywords : str
Keywords describing the topic to analyze (e.g. "transformers attention").
conferences : list of str, optional
Conference names to analyze (e.g. ["NeurIPS", "ICLR"]).
When omitted, the default conferences from the current session are used.
start_year : int, optional
Start year for analysis (inclusive).
end_year : int, optional
End year for analysis (inclusive).
"""
kwargs: Dict[str, Any] = {"topic_keywords": topic_keywords}
# Determine conferences: explicit param > deps default
if conferences:
kwargs["conferences"] = conferences
elif ctx.deps.conferences:
kwargs["conferences"] = ctx.deps.conferences
if start_year is not None:
kwargs["start_year"] = start_year
if end_year is not None:
kwargs["end_year"] = end_year
return _run_rag_tool(ctx, "get_topic_evolution", mcp_get_topic_evolution, kwargs)
def _tool_analyze_topic_relevance(
ctx: RunContext[RAGDeps],
topic: str,
distance_threshold: float = 1.1,
conference: Optional[str] = None,
years: Optional[List[int]] = None,
) -> str:
"""Analyze the relevance or popularity of a topic at a conference.
Use this tool to count how many papers are semantically similar to a
topic, measuring how prevalent or important a research topic is.
Returns relevance score, paper counts, and conference/year breakdowns.
A conference must be specified. If not provided by the user, the
currently selected conference from the web UI is used automatically.
The currently selected year is also applied automatically when not
explicitly provided.
Parameters
----------
ctx : RunContext[RAGDeps]
Agent context with dependencies.
topic : str
The topic or research question to analyze.
distance_threshold : float
Maximum distance to consider papers relevant (default: 1.1).
conference : str, optional
Conference name to analyze (e.g. "NeurIPS", "ICLR").
When omitted, the default conference from the current session is used.
years : list of int, optional
Filter results to specific years.
When omitted, the default year(s) from the current session are used.
"""
kwargs: Dict[str, Any] = {"topic": topic, "distance_threshold": distance_threshold}
# Determine conferences: explicit param > deps default
if conference:
kwargs["conferences"] = [conference]
elif ctx.deps.conferences:
kwargs["conferences"] = ctx.deps.conferences
# Determine years: explicit param > deps default
if years is not None:
kwargs["years"] = years
elif ctx.deps.years:
kwargs["years"] = ctx.deps.years
return _run_rag_tool(ctx, "analyze_topic_relevance", mcp_analyze_topic_relevance, kwargs)
def _tool_get_cluster_visualization(
ctx: RunContext[RAGDeps],
conference: Optional[str] = None,
) -> str:
"""Generate 2D visualization data for pre-computed clustered paper embeddings.
Use this tool to produce data for plotting papers grouped by topic
cluster. Returns points with x/y coordinates and cluster assignments.
A conference must be specified. If not provided by the user, the
currently selected conference from the web UI is used automatically.
Parameters
----------
ctx : RunContext[RAGDeps]
Agent context with dependencies.
conference : str, optional
Conference name to retrieve visualization for (e.g. "NeurIPS").
When omitted, the default conference from the current session is used.
"""
# Determine conferences: explicit param > deps default
if conference:
conferences = [conference]
elif ctx.deps.conferences:
conferences = ctx.deps.conferences
else:
conferences = None
years = ctx.deps.years if ctx.deps.years else None
kwargs = {"conferences": conferences, "years": years}
return _run_rag_tool(ctx, "get_cluster_visualization", mcp_get_cluster_visualization, kwargs)
def _tool_get_paper_details(
ctx: RunContext[RAGDeps],
title: Optional[str] = None,
paper_id: Optional[str] = None,
conference: Optional[str] = None,
year: Optional[int] = None,
limit: int = 5,
) -> str:
"""Get detailed metadata for specific papers from the database.
Use this tool when the user asks about: who wrote a paper, paper authors,
where to find a paper, PDF or poster links, session or room details,
paper awards, or any other metadata about a specific paper.
At least one of title or paper_id must be provided. A conference filter
is applied automatically from the current web UI selection when not
explicitly provided.
Parameters
----------
ctx : RunContext[RAGDeps]
Agent context with dependencies.
title : str, optional
Title or partial title to search for (case-insensitive).
paper_id : str, optional
Unique paper identifier (uid or original conference/OpenReview ID).
conference : str, optional
Filter by conference name (e.g. "NeurIPS", "ICLR").
When omitted, the default conference from the current session is used.
year : int, optional
Filter by publication year.
When omitted, the default year from the current session is used.
limit : int
Maximum number of papers to return when searching by title (default: 5).
"""
kwargs: Dict[str, Any] = {"limit": limit}
if title:
kwargs["title"] = title
if paper_id:
kwargs["paper_id"] = paper_id
# Determine conference: explicit param > deps default
if conference:
kwargs["conference"] = conference
elif ctx.deps.conferences:
kwargs["conference"] = ctx.deps.conferences[0]
# Determine year: explicit param > deps default
if year is not None:
kwargs["year"] = year
elif ctx.deps.years:
kwargs["year"] = ctx.deps.years[0]
return _run_rag_tool(ctx, "get_paper_details", mcp_get_paper_details, kwargs)
def _tool_summarize_interesting_papers(
ctx: RunContext[RAGDeps],
conference: Optional[str] = None,
years: Optional[List[int]] = None,
) -> str:
"""Summarize the papers the user marked as interesting.
Use this tool when the user asks for a summary, overview, or recap of the
papers they starred, liked, marked, or rated as interesting. The user's
ratings are provided automatically by the web UI, so do not ask for paper
IDs.
Parameters
----------
ctx : RunContext[RAGDeps]
Agent context with dependencies.
conference : str, optional
Conference filter (e.g. "ICLR" or "ICLR 2026"). When omitted, the
currently selected conference from the web UI is used automatically.
years : list of int, optional
Year filter. When omitted, the currently selected year from the web UI
is used automatically.
"""
priorities = ctx.deps.paper_priorities or {}
if not priorities:
return (
"No papers have been marked as interesting yet. " "Rate papers with the star buttons in the web UI first."
)
conferences = [conference] if conference else (ctx.deps.conferences or None)
resolved_years = years if years is not None else (ctx.deps.years or None)
kwargs: Dict[str, Any] = {"priorities": priorities}
if conferences:
kwargs["conference"] = conferences[0]
if resolved_years:
kwargs["years"] = resolved_years
return _run_rag_tool(ctx, "summarize_interesting_papers", mcp_summarize_interesting_papers, kwargs)
def _tool_recommend_related_papers(
ctx: RunContext[RAGDeps],
conferences: Optional[List[str]] = None,
years: Optional[List[int]] = None,
n_results: int = 10,
) -> str:
"""Recommend further papers based on the user's interesting papers.
Use this tool when the user asks for suggestions, recommendations, or
related papers based on the papers they starred, liked, marked, or rated
as interesting. Already-rated papers are excluded. The user's ratings are
provided automatically by the web UI, so do not ask for paper IDs.
Parameters
----------
ctx : RunContext[RAGDeps]
Agent context with dependencies.
conferences : list of str, optional
Restrict recommendations to these conferences. When omitted, the
currently selected conferences from the web UI are used automatically.
years : list of int, optional
Restrict recommendations to these years. When omitted, the currently
selected years from the web UI are used automatically.
n_results : int, optional
Number of recommendations to return (default: 10).
"""
priorities = ctx.deps.paper_priorities or {}
if not priorities:
return (
"No papers have been marked as interesting yet. " "Rate papers with the star buttons in the web UI first."
)
resolved_conferences = conferences if conferences else (ctx.deps.conferences or None)
resolved_years = years if years is not None else (ctx.deps.years or None)
kwargs: Dict[str, Any] = {"priorities": priorities, "n_results": n_results}
if resolved_conferences:
kwargs["conferences"] = resolved_conferences
if resolved_years:
kwargs["years"] = resolved_years
return _run_rag_tool(ctx, "recommend_related_papers", mcp_recommend_related_papers, kwargs)
[docs]
class RAGChat:
"""
RAG chat interface for querying conference papers.
Uses Pydantic AI as the agent framework to orchestrate tool calling and
response generation via an OpenAI-compatible language model API.
Parameters
----------
embeddings_manager : EmbeddingsManager
Manager for embeddings and vector search.
database : DatabaseManager
Database instance for querying paper details. REQUIRED.
lm_studio_url : str, optional
URL for OpenAI-compatible API, by default from config.
model : str, optional
Name of the language model, by default from config.
max_context_papers : int, optional
Maximum number of papers to include in context, by default from config.
temperature : float, optional
Sampling temperature for generation, by default from config.
enable_mcp_tools : bool, optional
Enable MCP clustering tools for topic analysis, by default True.
Examples
--------
>>> from abstracts_explorer.embeddings import EmbeddingsManager
>>> from abstracts_explorer.database import DatabaseManager
>>> em = EmbeddingsManager()
>>> em.init()
>>> db = DatabaseManager()
>>> chat = RAGChat(em, db)
>>>
>>> # Ask about specific papers
>>> response = chat.query("What are the latest advances in deep learning?")
>>> print(response)
>>>
>>> # Ask about conference topics (uses MCP tools automatically)
>>> response = chat.query("What are the main research topics at NeurIPS?")
>>> print(response)
"""
[docs]
def __init__(
self,
embeddings_manager,
database,
lm_studio_url: Optional[str] = None,
model: Optional[str] = None,
max_context_papers: Optional[int] = None,
temperature: Optional[float] = None,
enable_mcp_tools: bool = True,
):
if embeddings_manager is None:
raise RAGError("embeddings_manager is required.")
if database is None:
raise RAGError("database is required. RAGChat cannot operate without database access.")
config = get_config()
self.embeddings_manager = embeddings_manager
self.database = database
self.lm_studio_url = (lm_studio_url or config.llm_backend_url).rstrip("/")
self.model = model or config.chat_model
self.max_context_papers = max_context_papers or config.max_context_papers
self.temperature = temperature or config.chat_temperature
self.enable_mcp_tools = enable_mcp_tools
# Message history for multi-turn conversation
self._message_history: List[ModelMessage] = []
# Build the Pydantic AI agent
self.agent = self._build_agent()
def _build_base_instructions(self) -> str:
"""
Build the base system instructions for the RAG agent.
Returns
-------
str
Base instruction string without conference-specific context.
"""
return (
"You are an AI assistant helping researchers analyze conference data. "
"Use the available tools to search for papers, analyze topics, and understand trends. "
"Present the information in a clear, easy-to-understand format. "
"When the user asks about papers they marked, starred, liked, or rated as interesting, "
"use the summarize_interesting_papers or recommend_related_papers tools; their ratings "
"are provided automatically by the web UI, so do not ask for paper IDs. "
f"Today's date is {datetime.now().strftime('%Y-%m-%d')}. "
"When referencing specific papers, cite them using local markdown links: "
"[Paper-1](#paper-1), [Paper-2](#paper-2), etc."
)
def _build_instructions(
self,
conferences: Optional[List[str]] = None,
years: Optional[List[int]] = None,
available_conferences: Optional[List[str]] = None,
) -> str:
"""
Build context-aware system instructions including conference information.
Parameters
----------
conferences : list of str, optional
Currently selected conference(s) from the web UI.
years : list of int, optional
Currently selected year(s) from the web UI.
available_conferences : list of str, optional
All available conference names from registered plugins.
Returns
-------
str
Full instruction string with conference context appended when available.
"""
instructions = self._build_base_instructions()
context_parts = []
if conferences and years:
conf_str = ", ".join(conferences)
year_str = ", ".join(str(y) for y in years)
context_parts.append(
f"Unless the user specifies otherwise, assume the default conference is "
f"{conf_str} and the default year is {year_str}."
)
elif conferences:
conf_str = ", ".join(conferences)
context_parts.append(f"Unless the user specifies otherwise, assume the default conference is {conf_str}.")
elif years:
year_str = ", ".join(str(y) for y in years)
context_parts.append(f"Unless the user specifies otherwise, assume the default year is {year_str}.")
if available_conferences:
avail_str = ", ".join(available_conferences)
context_parts.append(f"The available conferences are: {avail_str}.")
if context_parts:
instructions += " " + " ".join(context_parts)
return instructions
def _build_agent(self) -> Agent[RAGDeps, str]:
"""
Build the Pydantic AI agent with tools and model configuration.
Returns
-------
Agent[RAGDeps, str]
Configured Pydantic AI agent.
"""
config = get_config()
requests_per_minute = config.requests_per_minute
# Build rate-limited async HTTP client when RPM > 0.
http_client: Optional[httpx2.AsyncClient] = None
if requests_per_minute > 0:
from abstracts_explorer.embeddings import RateLimitedTransport
transport = RateLimitedTransport(httpx2.AsyncHTTPTransport(), requests_per_minute=requests_per_minute)
http_client = httpx2.AsyncClient(transport=transport)
# Create OpenAI-compatible model
provider = OpenAIProvider(
base_url=f"{self.lm_studio_url}/v1",
api_key=config.llm_backend_auth_token or "lm-studio-local",
http_client=http_client,
)
ai_model = OpenAIChatModel(self.model, provider=provider)
# Build tool list
tools: List[Tool[RAGDeps]] = [
Tool(_tool_search_papers, takes_ctx=True, name="search_papers"),
Tool(_tool_get_paper_details, takes_ctx=True, name="get_paper_details"),
]
if self.enable_mcp_tools:
tools.extend(
[
Tool(_tool_get_conference_topics, takes_ctx=True, name="get_conference_topics"),
Tool(_tool_get_topic_evolution, takes_ctx=True, name="get_topic_evolution"),
Tool(_tool_analyze_topic_relevance, takes_ctx=True, name="analyze_topic_relevance"),
Tool(_tool_get_cluster_visualization, takes_ctx=True, name="get_cluster_visualization"),
Tool(
_tool_summarize_interesting_papers,
takes_ctx=True,
name="summarize_interesting_papers",
),
Tool(
_tool_recommend_related_papers,
takes_ctx=True,
name="recommend_related_papers",
),
]
)
return Agent(
ai_model,
deps_type=RAGDeps,
tools=tools,
)
[docs]
def query(
self,
question: str,
n_results: Optional[int] = None,
metadata_filter: Optional[Dict[str, Any]] = None,
system_prompt: Optional[str] = None,
conferences: Optional[List[str]] = None,
years: Optional[List[int]] = None,
available_conferences: Optional[List[str]] = None,
paper_priorities: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
"""
Query the RAG system with a question.
Parameters
----------
question : str
User's question.
n_results : int, optional
Number of papers to retrieve for context.
metadata_filter : dict, optional
Metadata filter for paper search (currently unused with Pydantic AI agent).
system_prompt : str, optional
Custom system prompt for the model (overrides default instructions including
conference context).
conferences : list of str, optional
Default conferences for MCP clustering tools. Typically set from
the web UI conference selector so that clustering tools use cached
results for the selected conference. Also included in the system prompt
so the LLM knows the current conference context.
years : list of int, optional
Default years for MCP clustering tools. Typically set from the
web UI year selector. Also included in the system prompt so the LLM
knows the current year context.
available_conferences : list of str, optional
All conference names available in the system (from registered plugins).
Included in the system prompt so the LLM can inform the user about
which conferences can be queried.
paper_priorities : dict, optional
Paper ratings from the web UI (``localStorage.paperPriorities``).
Used by the interesting-paper summary and recommendation tools.
Returns
-------
dict
Dictionary containing:
- response: str - Generated response
- papers: list - Retrieved papers used as context
- visualizations: list - Visualization data from tools
- metadata: dict - Additional metadata
Raises
------
RAGError
If query fails.
Examples
--------
>>> response = chat.query("What is attention mechanism?")
>>> print(response["response"])
>>> print(f"Based on {len(response['papers'])} papers")
"""
try:
if n_results is None:
n_results = self.max_context_papers
# Fresh deps for each query to capture tool results
deps = RAGDeps(
conferences=conferences or [],
years=years or [],
n_results=n_results,
paper_priorities=paper_priorities or {},
)
# Model settings
settings = OpenAIChatModelSettings(
temperature=self.temperature,
max_tokens=1000,
openai_reasoning_effort="low",
)
# Build instructions first so we can update message history
if system_prompt is not None:
instructions = system_prompt
else:
instructions = self._build_instructions(
conferences=conferences,
years=years,
available_conferences=available_conferences,
)
# Prepare run kwargs
run_kwargs: Dict[str, Any] = {
"deps": deps,
"model_settings": settings,
"instructions": instructions,
}
# Pass message history for conversation context, updating stale
# system prompt in the first request to reflect current conference.
# Pydantic AI doesn't re-add system prompts on continuation turns,
# so the first ModelRequest in history may contain outdated
# conference context from a prior turn.
if self._message_history:
self._update_message_history_instructions(instructions)
run_kwargs["message_history"] = self._message_history
# Run the agent
result = self.agent.run_sync(question, **run_kwargs)
# Update message history for next turn
self._message_history = result.all_messages()
# Extract papers from tool results
papers = self._extract_papers(deps.tool_results)
# Extract visualization data from tool results
visualizations = self._extract_visualizations(deps.tool_results)
# Determine which tools were executed
tools_executed = [tr["name"] for tr in deps.tool_results]
return {
"response": result.output,
"papers": papers,
"visualizations": visualizations,
"metadata": {
"n_papers": len(papers),
"model": self.model,
"used_tools": len(tools_executed) > 0,
"tools_executed": tools_executed,
},
}
except Exception as e:
raise RAGError(f"Query failed: {str(e)}") from e
[docs]
def chat(
self,
message: str,
use_context: bool = True,
n_results: Optional[int] = None,
) -> Dict[str, Any]:
"""
Continue conversation with context awareness.
Parameters
----------
message : str
User's message.
use_context : bool, optional
Whether to retrieve papers as context, by default True.
n_results : int, optional
Number of papers to retrieve.
Returns
-------
dict
Dictionary containing response and metadata.
Examples
--------
>>> response = chat.chat("Tell me more about transformers")
>>> print(response["response"])
"""
if use_context:
return self.query(message, n_results=n_results)
else:
# Use only conversation history without tool calls
try:
settings = OpenAIChatModelSettings(
temperature=self.temperature,
max_tokens=1000,
openai_reasoning_effort="low",
)
base_instructions = self._build_base_instructions()
run_kwargs: Dict[str, Any] = {
"deps": RAGDeps(),
"model_settings": settings,
"instructions": base_instructions,
}
if self._message_history:
self._update_message_history_instructions(base_instructions)
run_kwargs["message_history"] = self._message_history
result = self.agent.run_sync(message, **run_kwargs)
self._message_history = result.all_messages()
return {
"response": result.output,
"papers": [],
"metadata": {"n_papers": 0, "model": self.model},
}
except Exception as e:
raise RAGError(f"Chat failed: {str(e)}") from e
def _update_message_history_instructions(self, instructions: str) -> None:
"""
Update the instructions on the first ModelRequest in message history.
Pydantic AI does not re-add system prompts on continuation turns.
When the user changes the conference in the UI, the instructions need to
reflect the new conference, but the first ModelRequest in
``_message_history`` still holds the stale system prompt from the prior
conference. This method walks the history and updates the ``instructions``
field on every ModelRequest so the LLM sees consistent context.
"""
for msg in self._message_history:
if isinstance(msg, ModelRequest):
msg.instructions = instructions
[docs]
def reset_conversation(self):
"""
Reset conversation history.
Examples
--------
>>> chat.reset_conversation()
"""
self._message_history = []
logger.info("Conversation history reset")
@property
def conversation_history(self) -> List[Dict[str, str]]:
"""
Get conversation history as a list of role/content dicts.
This property converts Pydantic AI message objects to simple
dicts for backwards compatibility with code that accesses
``conversation_history`` directly.
Returns
-------
list of dict
Messages in ``{"role": "user"|"assistant", "content": "..."}`` format.
"""
history: List[Dict[str, str]] = []
for msg in self._message_history:
kind = getattr(msg, "kind", None)
if kind == "request":
for part in msg.parts:
part_kind = getattr(part, "part_kind", None)
if part_kind == "user-prompt":
content = getattr(part, "content", None)
if isinstance(content, str):
history.append({"role": "user", "content": content})
elif kind == "response":
for part in msg.parts:
part_kind = getattr(part, "part_kind", None)
if part_kind == "text":
content = getattr(part, "content", None)
if isinstance(content, str):
history.append({"role": "assistant", "content": content})
return history
[docs]
def export_conversation(self, output_path: Path) -> None:
"""
Export conversation history to JSON file.
Parameters
----------
output_path : Path
Path to output JSON file.
Examples
--------
>>> chat.export_conversation("conversation.json")
"""
output_path = Path(output_path)
history = self.conversation_history
with open(output_path, "w") as f:
json.dump(history, f, indent=2)
logger.info(f"Conversation exported to {output_path}")
@staticmethod
def _extract_papers(tool_results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""
Extract papers from tool results.
Parameters
----------
tool_results : list of dict
Tool results with 'name' and 'raw_result' keys.
Returns
-------
list of dict
Extracted papers.
"""
papers: List[Dict[str, Any]] = []
seen_uids: set[str] = set()
for tr in tool_results:
if tr["name"] in (
"search_papers",
"analyze_topic_relevance",
"summarize_interesting_papers",
"recommend_related_papers",
):
try:
result_json = json.loads(tr["raw_result"])
if "papers" in result_json:
for paper in result_json.get("papers", []):
uid = str(paper.get("uid", ""))
if uid and uid in seen_uids:
continue
papers.append(paper)
if uid:
seen_uids.add(uid)
except json.JSONDecodeError:
logger.warning(f"Failed to parse papers from {tr['name']} result")
return papers
@staticmethod
def _extract_visualizations(tool_results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""
Extract visualization data from tool results.
Converts raw MCP tool output from ``get_topic_evolution`` and
``get_cluster_visualization`` into lightweight chart descriptors
that the frontend can render with Plotly.
Parameters
----------
tool_results : list of dict
Tool results with 'name' and 'raw_result' keys.
Returns
-------
list of dict
Visualization descriptors. Each dict contains a ``type`` key
(``"topic_evolution"`` or ``"cluster_visualization"``) and the
data needed to draw the chart.
"""
visualizations: List[Dict[str, Any]] = []
topic_evolution_viz: Optional[Dict[str, Any]] = None
for tr in tool_results:
try:
data = json.loads(tr["raw_result"])
except (json.JSONDecodeError, TypeError, KeyError):
continue
if not isinstance(data, dict):
continue
if "error" in data:
continue
if tr["name"] == "get_topic_evolution":
conference_data = data.get("conference_data", {})
topic_name = data.get("topic", "")
if conference_data and topic_name:
# Merge into existing topic_evolution visualization if one exists
if topic_evolution_viz is not None:
topic_evolution_viz["topics"].append(topic_name)
topic_evolution_viz["conference_data"][topic_name] = conference_data
else:
topic_evolution_viz = {
"type": "topic_evolution",
"topics": [topic_name],
"conference_data": {topic_name: conference_data},
}
visualizations.append(topic_evolution_viz)
elif tr["name"] == "get_cluster_visualization":
points = data.get("points", [])
if points:
visualizations.append(
{
"type": "cluster_visualization",
"points": points,
"statistics": data.get("statistics", {}),
}
)
return visualizations