study-buddy / app /agents /wiki_agent.py
GitHub Actions
deploy d092bea3608b7a29952f16357fda39b7a29e399b
2e818da
Raw
History Blame Contribute Delete
6.71 kB
"""Wiki Agent -> streams grounded context cards for Infinite Wiki drill-downs.
Card format (streamed Markdown):
## [Term]
**What it means in this paper:** ...
**Why it matters here:** ...
**Evidence from the paper:** ...
**Good question to ask next:** ...
"""
from __future__ import annotations
import os
from app.agents.cerebras_client import CerebrasClient
_WIKI_SYSTEM_CONTENT_ONLY = """You are ResearchMate's Wiki bench inside a personal AI research lab.
Explain the selected paper term or passage using ONLY the provided project source material.
Do not invent facts. If the source material is thin, say what can be inferred from the source and what remains unclear.
Write a compact research note, not a classroom lesson.
Structure your response exactly as:
## {term}
**What it means in this paper:** 2-4 sentences grounded in the uploaded material.
**Why it matters here:**
* One project-specific reason.
* One connection to the surrounding passage or paper argument.
**Evidence from the paper:**
* A source-grounded fact with citation.
* A source-grounded fact with citation.
* A source-grounded formula, mechanism, or limitation if present.
**Good question to ask next:** One question that would help the reader understand or challenge the paper.
Cite inline with [Source: X, chunk N]. Ground every claim. Be concise."""
_WIKI_SYSTEM_NET_SUPPORT = """You are ResearchMate's Wiki bench inside a personal AI research lab.
Explain the selected paper term or passage. Use project source material first, then web source material only when the project source is thin or missing relevant context.
Do not invent facts. If grounding is thin, say what is grounded and what remains unclear.
Write a compact research note, not a classroom lesson.
First, check if the term/passage is discussed in the provided Source material.
- If it IS discussed in the Source material, explain it using the Source material, citing inline with [Source: X, chunk N].
- If it IS NOT discussed or lacks key details in the Source material, use the Web Source material to explain it. Use markdown links for web citations, formatted exactly as: [Source Title](url).
If both contain useful info, synthesize them, keeping citations clear.
Structure your response exactly as:
## {term}
**What it means in this paper:** 2-4 sentences grounded in the uploaded material, with web context only when needed.
**Why it matters here:**
* One project-specific reason.
* One connection to the surrounding passage or paper argument.
**Evidence from the paper and web:**
* A source-grounded fact with citation.
* A source-grounded fact or web-supported context with citation.
* A source-grounded formula, mechanism, or limitation if present.
**Good question to ask next:** One question that would help the reader understand or challenge the paper.
Keep citations accurate. Ground every claim. Be concise."""
class WikiAgent:
def __init__(self) -> None:
self._client = CerebrasClient()
async def search_tavily(self, query: str) -> list[dict]:
api_key = os.getenv("TAVILY_API_KEY", "")
if not api_key:
return []
import httpx
try:
async with httpx.AsyncClient() as client:
resp = await client.post(
"https://api.tavily.com/search",
json={
"api_key": api_key,
"query": query,
"max_results": 3
},
timeout=5.0
)
if resp.status_code == 200:
return resp.json().get("results", [])
except Exception as e:
print("Tavily search error:", e)
return []
async def stream_card(
self,
selection_text: str,
surrounding_context: str,
chunks: list[dict],
familiarity: str,
parent_context: str = "",
knowledge_mode: str = "content_only",
project_memory: str = "",
student_memory: str = "",
):
"""Async generator -> yields text tokens for the Infinite Wiki card."""
chunk_text = "\n\n".join(
f"[Source: {c.get('source', '?')}]\n{c['text']}"
for i, c in enumerate(chunks)
)
parent_note = f"\nDrill-down context: {parent_context[:300]}" if parent_context else ""
system_prompt = _WIKI_SYSTEM_NET_SUPPORT if knowledge_mode == "net_support" else _WIKI_SYSTEM_CONTENT_ONLY
# Fetch optional external context. Project memory is intentionally not queried
# here until it can be scoped by project_id; the old global_curriculum dataset
# was not guaranteed to exist and produced noisy, misleading warnings.
tavily_results = await self.search_tavily(selection_text) if knowledge_mode == "net_support" else []
web_text = ""
if tavily_results:
web_text = "\n\n".join(
f"[Web Source: {r.get('title')}, URL: {r.get('url')}]\n{r.get('content')}"
for r in tavily_results
)
# Resolve complexity level names dynamically
level_name = {
"eli5": "5-year old",
"high_school": "high school",
"graduate": "graduate university",
"expert": "expert research specialist"
}.get(familiarity, "high school")
# Inject actual selection term and level variables into the system prompt structure.
# Level affects depth/tone, not the visible card headings.
system_prompt = system_prompt.replace("{term}", selection_text)
system_prompt = system_prompt.replace("{level}", level_name)
user_content = (
f"Term/passage to explain: \"{selection_text}\"\n"
f"Surrounding text: {surrounding_context[:400]}{parent_note}\n\n"
f"Familiarity level: {familiarity}\n\n"
f"Source material:\n{chunk_text}"
)
if project_memory:
user_content += (
"\n\nProject memory (research state, not source evidence; never cite as a paper fact):\n"
f"{project_memory}"
)
if student_memory:
user_content += (
"\n\nStudent memory (adapt tone/depth only; never cite as a paper fact):\n"
f"{student_memory}"
)
if web_text:
user_content += f"\n\nWeb Source material:\n{web_text}"
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_content},
]
async for token in self._client.stream_complete(messages):
yield token