File size: 6,711 Bytes
2e818da
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
"""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