44
55import importlib
66import time
7- from typing import Any , Dict , List
7+ from typing import Any , Dict , Iterable , Iterator , List , Optional
88
99from paperflow .providers import build_llm_provider
1010
1515SYSTEM_PROMPT = """You are PaperFlow's private research wiki assistant.
1616
1717Answer only from the provided wiki snippets. Cite every concrete claim with
18- [N], where N is the snippet number. If the snippets do not contain enough
19- evidence, say that the local wiki does not have enough material yet.
20- Use concise Chinese by default unless the user asks otherwise."""
18+ [N], where N is the snippet number shown in the context. Keep citation markers
19+ inline immediately after the supported claim, and do not invent references. If
20+ the snippets do not contain enough evidence, say that the local wiki does not
21+ have enough material yet. Do not expose internal section labels such as Q1/Q2
22+ unless the user explicitly asks for report-section structure. Use concise
23+ Chinese by default unless the user asks otherwise."""
2124
2225
2326def _snippet (node : Dict [str , Any ], max_chars : int = 700 ) -> str :
@@ -51,7 +54,7 @@ def _build_prompt(question: str, hits: List[Dict[str, Any]]) -> str:
5154
5255def _extractive_fallback (hits : List [Dict [str , Any ]], error : str ) -> str :
5356 lines = [
54- "The configured LLM was not available, so PaperFlow returned local wiki snippets instead. " ,
57+ "当前 LLM 不可用， PaperFlow 先返回本地 Wiki 证据片段（ local wiki snippets）。 " ,
5558 f"Provider error: { error } ,
5659 "" ,
5760 ]
@@ -62,65 +65,175 @@ def _extractive_fallback(hits: List[Dict[str, Any]], error: str) -> str:
6265 return "\n " .join (lines ).strip ()
6366
6467
65- def answer_question (user_id : str , question : str , * , limit : int = 8 ) -> Dict [str , Any ]:
66- """Return an LLM answer with local wiki citations."""
67- started = time .time ()
68+ def _merge_hits (
69+ pinned_nodes : Iterable [Dict [str , Any ]],
70+ retrieved_nodes : Iterable [Dict [str , Any ]],
71+ limit : int ,
72+ ) -> List [Dict [str , Any ]]:
73+ merged : List [Dict [str , Any ]] = []
74+ seen = set ()
75+ for node in [* list (pinned_nodes or []), * list (retrieved_nodes or [])]:
76+ node_id = str (node .get ("node_id" ) or "" ).strip ()
77+ if not node_id or node_id in seen :
78+ continue
79+ seen .add (node_id )
80+ merged .append (node )
81+ if len (merged ) >= max (1 , int (limit )):
82+ break
83+ return merged
84+
85+
86+ def _prepare_hits (
87+ user_id : str ,
88+ question : str ,
89+ limit : int ,
90+ pinned_nodes : Optional [List [Dict [str , Any ]]],
91+ ) -> tuple [List [Dict [str , Any ]], Optional [str ]]:
6892 embedding_error = None
6993 try :
7094 wiki_db .embed_nodes_for_user (user_id , limit = 500 )
7195 except Exception as exc :
7296 embedding_error = str (exc )
73- hits = wiki_db .search_nodes (user_id , question , limit = limit )
74- if not hits :
75- return {
76- "text" : (
77- "The local wiki does not have enough relevant material yet. "
78- "Run `paperflow read` for related papers, or run "
79- "`paperflow wiki backfill` to import existing runtime history."
80- ),
81- "citations" : [],
82- "elapsed_ms" : int ((time .time () - started ) * 1000 ),
83- "token_usage" : {},
84- }
97+ retrieved = wiki_db .search_nodes (user_id , question , limit = max (limit , limit + len (pinned_nodes or [])))
98+ return _merge_hits (pinned_nodes or [], retrieved , limit ), embedding_error
8599
86- prompt = _build_prompt (question , hits )
87- llm = build_llm_provider ()
88- llm_error = None
89- response = None
90- try :
91- response = llm .generate (prompt , system = SYSTEM_PROMPT , temperature = 0.0 , max_tokens = 900 )
92- answer_text = response .text
93- except Exception as exc :
94- llm_error = str (exc )
95- answer_text = _extractive_fallback (hits , llm_error )
96- citations_by_node = wiki_db .get_citations_for_nodes (user_id , [hit ["node_id" ] for hit in hits ])
100+
101+ def _build_citations (user_id : str , hits : List [Dict [str , Any ]]) -> List [Dict [str , Any ]]:
102+ citations_by_node = wiki_db .get_citations_for_nodes (user_id , [hit ["node_id" ] for hit in hits if hit .get ("node_id" )])
97103 citations = []
98104 for index , hit in enumerate (hits , start = 1 ):
105+ metadata = hit .get ("metadata" ) or {}
106+ display_node = hit
107+ section_title = ""
108+ if str (hit .get ("node_type" ) or "" ) == "section" and metadata .get ("parent_paper_id" ):
109+ parent = wiki_db .get_node (user_id , str (metadata .get ("parent_paper_id" )))
110+ if parent :
111+ display_node = parent
112+ section_title = str (hit .get ("title" ) or "" )
99113 node_citations = citations_by_node .get (hit ["node_id" ], [])
100114 first = node_citations [0 ] if node_citations else {}
115+ display_metadata = dict (display_node .get ("metadata" ) or {})
116+ if metadata .get ("parent_paper_id" ):
117+ display_metadata .setdefault ("parent_paper_id" , metadata .get ("parent_paper_id" ))
118+ if section_title :
119+ display_metadata ["section_title" ] = section_title
120+ display_metadata ["evidence_node_id" ] = hit ["node_id" ]
101121 citations .append (
102122 {
103123 "index" : index ,
104- "node_id" : hit ["node_id" ],
105- "title" : hit ["title" ],
106- "node_type" : hit ["node_type" ],
124+ "node_id" : display_node ["node_id" ],
125+ "title" : display_node ["title" ],
126+ "node_type" : display_node ["node_type" ],
107127 "excerpt" : first .get ("excerpt" ) or _snippet (hit , max_chars = 260 ),
108128 "source_type" : first .get ("source" ),
109129 "source_id" : first .get ("source_id" ),
110130 "anchor" : first .get ("anchor" ),
111- "metadata" : hit . get ( "metadata" ) or {} ,
131+ "metadata" : display_metadata ,
112132 }
113133 )
134+ return citations
135+
136+
137+ def _empty_answer (started : float ) -> Dict [str , Any ]:
138+ return {
139+ "text" : (
140+ "本地 Wiki 还没有足够相关的材料。可以先对相关论文生成精读报告，"
141+ "或运行 `paperflow wiki backfill` 导入已有运行历史。"
142+ ),
143+ "citations" : [],
144+ "elapsed_ms" : int ((time .time () - started ) * 1000 ),
145+ "token_usage" : {},
146+ }
147+
148+
149+ def _token_usage (llm : Any , response : Any , embedding_error : Optional [str ], llm_error : Optional [str ]) -> Dict [str , Any ]:
150+ return {
151+ "provider" : getattr (llm , "name" , "unknown" ),
152+ "model" : getattr (llm , "model" , "unknown" ),
153+ "prompt_tokens" : response .prompt_tokens if response else 0 ,
154+ "completion_tokens" : response .completion_tokens if response else 0 ,
155+ "embedding_error" : embedding_error ,
156+ "llm_error" : llm_error ,
157+ }
158+
159+
160+ def answer_question (
161+ user_id : str ,
162+ question : str ,
163+ * ,
164+ limit : int = 8 ,
165+ pinned_nodes : Optional [List [Dict [str , Any ]]] = None ,
166+ ) -> Dict [str , Any ]:
167+ """Return an LLM answer with local wiki citations."""
168+ started = time .time ()
169+ hits , embedding_error = _prepare_hits (user_id , question , limit , pinned_nodes )
170+ if not hits :
171+ return _empty_answer (started )
172+
173+ prompt = _build_prompt (question , hits )
174+ llm = build_llm_provider ()
175+ llm_error = None
176+ response = None
177+ try :
178+ response = llm .generate (prompt , system = SYSTEM_PROMPT , temperature = 0.0 , max_tokens = 900 )
179+ answer_text = response .text
180+ except Exception as exc :
181+ llm_error = str (exc )
182+ answer_text = _extractive_fallback (hits , llm_error )
183+ citations = _build_citations (user_id , hits )
114184 return {
115185 "text" : answer_text ,
116186 "citations" : citations ,
117187 "elapsed_ms" : int ((time .time () - started ) * 1000 ),
118- "token_usage" : {
119- "provider" : getattr (llm , "name" , "unknown" ),
120- "model" : getattr (llm , "model" , "unknown" ),
121- "prompt_tokens" : response .prompt_tokens if response else 0 ,
122- "completion_tokens" : response .completion_tokens if response else 0 ,
123- "embedding_error" : embedding_error ,
124- "llm_error" : llm_error ,
125- },
188+ "token_usage" : _token_usage (llm , response , embedding_error , llm_error ),
189+ }
190+
191+
192+ def answer_question_stream (
193+ user_id : str ,
194+ question : str ,
195+ * ,
196+ limit : int = 8 ,
197+ pinned_nodes : Optional [List [Dict [str , Any ]]] = None ,
198+ ) -> Iterator [Dict [str , Any ]]:
199+ """Yield answer events with local citations and provider-native chunks."""
200+ started = time .time ()
201+ hits , embedding_error = _prepare_hits (user_id , question , limit , pinned_nodes )
202+ if not hits :
203+ result = _empty_answer (started )
204+ yield {"event" : "meta" , "data" : {key : value for key , value in result .items () if key != "text" }}
205+ yield {"event" : "chunk" , "data" : {"text" : result ["text" ]}}
206+ yield {"event" : "done" , "data" : result }
207+ return
208+
209+ citations = _build_citations (user_id , hits )
210+ llm = build_llm_provider ()
211+ prompt = _build_prompt (question , hits )
212+ llm_error = None
213+ text_parts : List [str ] = []
214+ meta = {
215+ "citations" : citations ,
216+ "elapsed_ms" : 0 ,
217+ "token_usage" : _token_usage (llm , None , embedding_error , None ),
218+ "streaming" : {"provider" : True , "transport" : "sse" },
219+ }
220+ yield {"event" : "meta" , "data" : meta }
221+ try :
222+ for chunk in llm .stream_generate (prompt , system = SYSTEM_PROMPT , temperature = 0.0 , max_tokens = 900 ):
223+ if not chunk :
224+ continue
225+ text_parts .append (chunk )
226+ yield {"event" : "chunk" , "data" : {"text" : chunk }}
227+ except Exception as exc :
228+ llm_error = str (exc )
229+ fallback = _extractive_fallback (hits , llm_error )
230+ text_parts = [fallback ]
231+ yield {"event" : "chunk" , "data" : {"text" : fallback }}
232+
233+ result = {
234+ "text" : "" .join (text_parts ),
235+ "citations" : citations ,
236+ "elapsed_ms" : int ((time .time () - started ) * 1000 ),
237+ "token_usage" : _token_usage (llm , None , embedding_error , llm_error ),
126238 }
239+ yield {"event" : "done" , "data" : result }
0 commit comments