Update project files

This commit is contained in:
Xin Wang
2026-03-30 23:44:48 +08:00
parent e34e569de4
commit 00ef163767
11 changed files with 137 additions and 261 deletions

View File

@@ -21,7 +21,7 @@ def main() -> None:
)
parser.add_argument(
"--reasoning-model",
default="gemini-2.5-pro-preview-05-06",
default="GLM-4.5-Air",
help="Model for the final answer",
)
args = parser.parse_args()

View File

@@ -11,13 +11,13 @@ requires-python = ">=3.11,<4.0"
dependencies = [
"langgraph>=0.2.6",
"langchain>=0.3.19",
"langchain-google-genai",
"langchain-openai",
"python-dotenv>=1.0.1",
"langgraph-sdk>=0.1.57",
"langgraph-cli",
"langgraph-api",
"fastapi",
"google-genai",
"tavily-python",
]

View File

@@ -9,21 +9,21 @@ class Configuration(BaseModel):
"""The configuration for the agent."""
query_generator_model: str = Field(
default="gemini-2.0-flash",
default="GLM-4.5-Air",
metadata={
"description": "The name of the language model to use for the agent's query generation."
},
)
reflection_model: str = Field(
default="gemini-2.5-flash",
default="GLM-4.5-Air",
metadata={
"description": "The name of the language model to use for the agent's reflection."
},
)
answer_model: str = Field(
default="gemini-2.5-pro",
default="GLM-4.5-Air",
metadata={
"description": "The name of the language model to use for the agent's answer."
},

View File

@@ -7,7 +7,8 @@ from langgraph.types import Send
from langgraph.graph import StateGraph
from langgraph.graph import START, END
from langchain_core.runnables import RunnableConfig
from google.genai import Client
from langchain_openai import ChatOpenAI
from tavily import TavilyClient
from agent.state import (
OverallState,
@@ -23,29 +24,38 @@ from agent.prompts import (
reflection_instructions,
answer_instructions,
)
from langchain_google_genai import ChatGoogleGenerativeAI
from agent.utils import (
get_citations,
deduplicate_and_format_sources,
get_research_topic,
insert_citation_markers,
resolve_urls,
)
load_dotenv()
if os.getenv("GEMINI_API_KEY") is None:
raise ValueError("GEMINI_API_KEY is not set")
if os.getenv("OPENAI_API_KEY") is None:
raise ValueError("OPENAI_API_KEY is not set")
# Used for Google Search API
genai_client = Client(api_key=os.getenv("GEMINI_API_KEY"))
if os.getenv("TAVILY_API_KEY") is None:
raise ValueError("TAVILY_API_KEY is not set")
tavily_client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
def create_llm(model: str, temperature: float) -> ChatOpenAI:
"""Create an OpenAI-compatible chat client."""
return ChatOpenAI(
model=model,
temperature=temperature,
max_retries=2,
api_key=os.getenv("OPENAI_API_KEY"),
base_url=os.getenv("OPENAI_BASE_URL", "https://open.bigmodel.cn/api/paas/v4"),
)
# Nodes
def generate_query(state: OverallState, config: RunnableConfig) -> QueryGenerationState:
"""LangGraph node that generates search queries based on the User's question.
Uses Gemini 2.0 Flash to create an optimized search queries for web research based on
the User's question.
Uses the configured chat model to create optimized search queries for web research.
Args:
state: Current graph state containing the User's question
@@ -56,27 +66,19 @@ def generate_query(state: OverallState, config: RunnableConfig) -> QueryGenerati
"""
configurable = Configuration.from_runnable_config(config)
# check for custom initial search query count
# Check for custom initial search query count.
if state.get("initial_search_query_count") is None:
state["initial_search_query_count"] = configurable.number_of_initial_queries
# init Gemini 2.0 Flash
llm = ChatGoogleGenerativeAI(
model=configurable.query_generator_model,
temperature=1.0,
max_retries=2,
api_key=os.getenv("GEMINI_API_KEY"),
)
llm = create_llm(configurable.query_generator_model, temperature=1.0)
structured_llm = llm.with_structured_output(SearchQueryList)
# Format the prompt
current_date = get_current_date()
formatted_prompt = query_writer_instructions.format(
current_date=current_date,
research_topic=get_research_topic(state["messages"]),
number_queries=state["initial_search_query_count"],
)
# Generate the search queries
result = structured_llm.invoke(formatted_prompt)
return {"search_query": result.query}
@@ -93,46 +95,34 @@ def continue_to_web_research(state: QueryGenerationState):
def web_research(state: WebSearchState, config: RunnableConfig) -> OverallState:
"""LangGraph node that performs web research using the native Google Search API tool.
Executes a web search using the native Google Search API tool in combination with Gemini 2.0 Flash.
Args:
state: Current graph state containing the search query and research loop count
config: Configuration for the runnable, including search API settings
Returns:
Dictionary with state update, including sources_gathered, research_loop_count, and web_research_results
"""
# Configure
configurable = Configuration.from_runnable_config(config)
"""LangGraph node that performs web research using Tavily search."""
Configuration.from_runnable_config(config)
formatted_prompt = web_searcher_instructions.format(
current_date=get_current_date(),
research_topic=state["search_query"],
)
search_response = tavily_client.search(
query=state["search_query"],
topic="general",
search_depth="advanced",
max_results=5,
include_answer=False,
include_raw_content=False,
)
sources_gathered, formatted_results = deduplicate_and_format_sources(
search_response.get("results", []), int(state["id"])
)
# Uses the google genai client as the langchain client doesn't return grounding metadata
response = genai_client.models.generate_content(
model=configurable.query_generator_model,
contents=formatted_prompt,
config={
"tools": [{"google_search": {}}],
"temperature": 0,
},
web_research_result = (
f"{formatted_prompt}\n\n"
f"Search query: {state['search_query']}\n\n"
f"Findings:\n{formatted_results}"
)
# resolve the urls to short urls for saving tokens and time
resolved_urls = resolve_urls(
response.candidates[0].grounding_metadata.grounding_chunks, state["id"]
)
# Gets the citations and adds them to the generated text
citations = get_citations(response, resolved_urls)
modified_text = insert_citation_markers(response.text, citations)
sources_gathered = [item for citation in citations for item in citation["segments"]]
return {
"sources_gathered": sources_gathered,
"search_query": [state["search_query"]],
"web_research_result": [modified_text],
"web_research_result": [web_research_result],
}
@@ -162,13 +152,7 @@ def reflection(state: OverallState, config: RunnableConfig) -> ReflectionState:
research_topic=get_research_topic(state["messages"]),
summaries="\n\n---\n\n".join(state["web_research_result"]),
)
# init Reasoning Model
llm = ChatGoogleGenerativeAI(
model=reasoning_model,
temperature=1.0,
max_retries=2,
api_key=os.getenv("GEMINI_API_KEY"),
)
llm = create_llm(reasoning_model, temperature=1.0)
result = llm.with_structured_output(Reflection).invoke(formatted_prompt)
return {
@@ -241,13 +225,7 @@ def finalize_answer(state: OverallState, config: RunnableConfig):
summaries="\n---\n\n".join(state["web_research_result"]),
)
# init Reasoning Model, default to Gemini 2.5 Flash
llm = ChatGoogleGenerativeAI(
model=reasoning_model,
temperature=0,
max_retries=2,
api_key=os.getenv("GEMINI_API_KEY"),
)
llm = create_llm(reasoning_model, temperature=0)
result = llm.invoke(formatted_prompt)
# Replace the short urls with the original urls and add all used urls to the sources_gathered

View File

@@ -34,11 +34,11 @@ Topic: What revenue grew more last year apple stock or the number of people buyi
Context: {research_topic}"""
web_searcher_instructions = """Conduct targeted Google Searches to gather the most recent, credible information on "{research_topic}" and synthesize it into a verifiable text artifact.
web_searcher_instructions = """Conduct targeted web research to gather the most recent, credible information on "{research_topic}" and synthesize it into a verifiable text artifact.
Instructions:
- Query should ensure that the most current information is gathered. The current date is {current_date}.
- Conduct multiple, diverse searches to gather comprehensive information.
- Review the returned web search results and extract the strongest, most relevant findings.
- Consolidate key findings while meticulously tracking the source(s) for each specific piece of information.
- The output should be a well-written summary or report based on your search findings.
- Only include the information found in the search results, don't make up any information.
@@ -87,7 +87,7 @@ Instructions:
- You have access to all the information gathered from the previous steps.
- You have access to the user's question.
- Generate a high-quality answer to the user's question based on the provided summaries and the user's question.
- Include the sources you used from the Summaries in the answer correctly, use markdown format (e.g. [apnews](https://vertexaisearch.cloud.google.com/id/1-0)). THIS IS A MUST.
- Include the sources you used from the Summaries in the answer correctly, using markdown links. THIS IS A MUST.
User Context:
- {research_topic}

View File

@@ -1,166 +1,59 @@
from typing import Any, Dict, List
from langchain_core.messages import AnyMessage, AIMessage, HumanMessage
from typing import Any
from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
def get_research_topic(messages: List[AnyMessage]) -> str:
"""
Get the research topic from the messages.
"""
# check if request has a history and combine the messages into a single string
def get_research_topic(messages: list[AnyMessage]) -> str:
"""Get the research topic from the messages."""
if len(messages) == 1:
research_topic = messages[-1].content
else:
research_topic = ""
for message in messages:
if isinstance(message, HumanMessage):
research_topic += f"User: {message.content}\n"
elif isinstance(message, AIMessage):
research_topic += f"Assistant: {message.content}\n"
return str(messages[-1].content)
research_topic = ""
for message in messages:
if isinstance(message, HumanMessage):
research_topic += f"User: {message.content}\n"
elif isinstance(message, AIMessage):
research_topic += f"Assistant: {message.content}\n"
return research_topic
def resolve_urls(urls_to_resolve: List[Any], id: int) -> Dict[str, str]:
"""
Create a map of the vertex ai search urls (very long) to a short url with a unique id for each url.
Ensures each original URL gets a consistent shortened form while maintaining uniqueness.
"""
prefix = f"https://vertexaisearch.cloud.google.com/id/"
urls = [site.web.uri for site in urls_to_resolve]
def deduplicate_and_format_sources(
results: list[dict[str, Any]], search_id: int
) -> tuple[list[dict[str, str]], str]:
"""Convert Tavily results into source metadata and a markdown research summary."""
sources: list[dict[str, str]] = []
formatted_sections: list[str] = []
seen_urls: set[str] = set()
# Create a dictionary that maps each unique URL to its first occurrence index
resolved_map = {}
for idx, url in enumerate(urls):
if url not in resolved_map:
resolved_map[url] = f"{prefix}{id}-{idx}"
for idx, result in enumerate(results):
url = result.get("url")
if not url or url in seen_urls:
continue
return resolved_map
seen_urls.add(url)
title = result.get("title") or f"Source {idx + 1}"
short_url = f"https://search.local/{search_id}-{idx}"
content = (result.get("content") or "").strip()
def insert_citation_markers(text, citations_list):
"""
Inserts citation markers into a text string based on start and end indices.
Args:
text (str): The original text string.
citations_list (list): A list of dictionaries, where each dictionary
contains 'start_index', 'end_index', and
'segment_string' (the marker to insert).
Indices are assumed to be for the original text.
Returns:
str: The text with citation markers inserted.
"""
# Sort citations by end_index in descending order.
# If end_index is the same, secondary sort by start_index descending.
# This ensures that insertions at the end of the string don't affect
# the indices of earlier parts of the string that still need to be processed.
sorted_citations = sorted(
citations_list, key=lambda c: (c["end_index"], c["start_index"]), reverse=True
)
modified_text = text
for citation_info in sorted_citations:
# These indices refer to positions in the *original* text,
# but since we iterate from the end, they remain valid for insertion
# relative to the parts of the string already processed.
end_idx = citation_info["end_index"]
marker_to_insert = ""
for segment in citation_info["segments"]:
marker_to_insert += f" [{segment['label']}]({segment['short_url']})"
# Insert the citation marker at the original end_idx position
modified_text = (
modified_text[:end_idx] + marker_to_insert + modified_text[end_idx:]
sources.append(
{
"label": title,
"short_url": short_url,
"value": url,
}
)
return modified_text
def get_citations(response, resolved_urls_map):
"""
Extracts and formats citation information from a Gemini model's response.
This function processes the grounding metadata provided in the response to
construct a list of citation objects. Each citation object includes the
start and end indices of the text segment it refers to, and a string
containing formatted markdown links to the supporting web chunks.
Args:
response: The response object from the Gemini model, expected to have
a structure including `candidates[0].grounding_metadata`.
It also relies on a `resolved_map` being available in its
scope to map chunk URIs to resolved URLs.
Returns:
list: A list of dictionaries, where each dictionary represents a citation
and has the following keys:
- "start_index" (int): The starting character index of the cited
segment in the original text. Defaults to 0
if not specified.
- "end_index" (int): The character index immediately after the
end of the cited segment (exclusive).
- "segments" (list[str]): A list of individual markdown-formatted
links for each grounding chunk.
- "segment_string" (str): A concatenated string of all markdown-
formatted links for the citation.
Returns an empty list if no valid candidates or grounding supports
are found, or if essential data is missing.
"""
citations = []
# Ensure response and necessary nested structures are present
if not response or not response.candidates:
return citations
candidate = response.candidates[0]
if (
not hasattr(candidate, "grounding_metadata")
or not candidate.grounding_metadata
or not hasattr(candidate.grounding_metadata, "grounding_supports")
):
return citations
for support in candidate.grounding_metadata.grounding_supports:
citation = {}
# Ensure segment information is present
if not hasattr(support, "segment") or support.segment is None:
continue # Skip this support if segment info is missing
start_index = (
support.segment.start_index
if support.segment.start_index is not None
else 0
formatted_sections.append(
"\n".join(
[
f"- {title} [{title}]({short_url})",
f" URL: {url}",
f" Snippet: {content or 'No snippet returned.'}",
]
)
)
# Ensure end_index is present to form a valid segment
if support.segment.end_index is None:
continue # Skip if end_index is missing, as it's crucial
if not formatted_sections:
formatted_sections.append("- No Tavily results were returned for this query.")
# Add 1 to end_index to make it an exclusive end for slicing/range purposes
# (assuming the API provides an inclusive end_index)
citation["start_index"] = start_index
citation["end_index"] = support.segment.end_index
citation["segments"] = []
if (
hasattr(support, "grounding_chunk_indices")
and support.grounding_chunk_indices
):
for ind in support.grounding_chunk_indices:
try:
chunk = candidate.grounding_metadata.grounding_chunks[ind]
resolved_url = resolved_urls_map.get(chunk.web.uri, None)
citation["segments"].append(
{
"label": chunk.web.title.split(".")[:-1][0],
"short_url": resolved_url,
"value": chunk.web.uri,
}
)
except (IndexError, AttributeError, NameError):
# Handle cases where chunk, web, uri, or resolved_map might be problematic
# For simplicity, we'll just skip adding this particular segment link
# In a production system, you might want to log this.
pass
citations.append(citation)
return citations
return sources, "\n".join(formatted_sections)