Updates to _core, improvements to research workflow

- Improve stability, performance
- Extend config_app.yaml, e.g. more models tested
- Refactor app UI texts to separate file
- Track token usage of last research step
- Log input to research writing to JSON files
- Bug fixes
- Resolve some minor issues in UI

Author: rnckp

.gitignore

@@ -172,5 +172,7 @@ cython_debug/
 
 # PyPI configuration file
 .pypirc
+
+# Additional files to ignore
 .DS_Store
 *.docx

01_data/_weaviate_data/classifications.db

@@ -0,0 +1,29 @@
+INFO_TEXT_MODAL = """Mit dieser App kannst du **vertieft über eigene Dokumentsammlungen recherchieren**.
+
+Die App dient zum Testen. **Beachte, dass Sprachmodelle (LLMs) Fehler machen und die Ergebnisse fehlerhaft oder unvollständig sein können.** Überprüfe die Ergebnisse immer.
+
+Deine Fragen werden an Clouddienste weitergeleitet und dort verarbeitet. **Gib daher nur als öffentlich klassifizierte Informationen als Fragen ein!** Beachte auch, dass die Nutzung anonymisiert aufgezeichnet wird und Mitarbeitende vom Statistischen Amt Eingaben stichprobenartig überprüfen, um die App zu verbessern.
+
+Zu Demonstrationszwecken bezieht die App für die Antworten eine kleine Auswahl von [Kantonsratsprotokollen des Kantons Zürich ein](https://opendata.swiss/de/dataset/zurcher-kantonsratsprotokolle-des-19-jahrhunderts). 
+
+Verantwortlich: Statistisches Amt, [Team Data](mailto:[email protected]).
+
+App-Version v0.2. Letzte Aktualisierung 23.7.2025
+
+### Wie funktioniert die App?
+
+Die App arbeitet in mehreren Schritten:
+
+1. **Suchanfragen formulieren**: Basierend auf deiner Fragestellung generiert die App gezielte Suchanfragen.  
+2. **Recherche durchführen**: Anschließend durchsucht die App die Dokumente nach passenden Textstellen.  
+3. **Relevanz prüfen**: Die gefundenen Passagen werden daraufhin geprüft, ob sie für deine Fragestellung von Bedeutung sind.  
+4. **Ganze Beschlüsse analysieren**: Zu relevanten Textstellen werden die Dokumente im Volltext analysiert und inhaltlich in Bezug auf die Frage zusammengefasst.  
+5. **Recherchestand reflektieren**: Falls die iterative Recherche aktiviert ist, bewertet die App den bisherigen Erkenntnisstand und entscheidet, ob weitere Recherchen notwendig sind.  
+6. **Abschlussbericht erstellen**: Die App fasst die Ergebnisse in einem Abschlussbericht zusammen.
+""".strip()
+
+INFO_TEXT_SIDEBAR = """
+Recherchewerkzeug für eigene Dokumentsammlungen.\n\n:red[Achtung: Dies ist ein experimenteller Prototyp. Gib nur als öffentlich klassifizierte Daten als Fragen ein. Die Ergebnisse können fehlerhaft oder unvollständig sein. **Überprüfe die Ergebnisse immer.**] \n\n Die Bearbeitung kann einige Minuten dauern, abhängig von der Komplexität der Anfrage und der Anzahl der relevanten Dokumente.
+""".strip()
+
+SAMPLE_QUERY = "Was hat der Kantonsrat zu Steuern entschieden?"

01_data/_weaviate_data/raft/raft.db

@@ -1,11 +1,15 @@
 from abc import ABC, abstractmethod
-from typing import Optional, Dict, Any, Iterator
+from typing import Optional, Dict, Any
 import os
+import requests
+import json
+from datetime import datetime
 from openai import OpenAI
 from tenacity import retry, stop_after_attempt, wait_random_exponential
 from dotenv import load_dotenv
 from _core.config import config
 from _core.logger import custom_logger
+from _core.utils import TokenCounter
 
 
 try:
@@ -32,8 +36,8 @@ def call_structured(
         pass
 
     @abstractmethod
-    def call_streamed(self, prompt: str, **kwargs) -> Iterator:
-        """Stream responses from the LLM."""
+    def call_with_reasoning(self, prompt: str, **kwargs) -> tuple[str, dict]:
+        """Call LLM API with reasoning parameters."""
         pass
 
 
@@ -78,7 +82,7 @@ def _call():
             completion = self.client.chat.completions.create(
                 model=model_id or config["models"]["performance_low"],
                 temperature=temperature or config["temperature"]["low"],
-                max_tokens=max_tokens or config["llm"]["max_tokens"],
+                max_tokens=max_tokens or config["llm"]["max_tokens_output"],
                 reasoning_effort=reasoning_effort,
                 messages=[{"role": "user", "content": prompt}],
                 **kwargs,
@@ -104,7 +108,7 @@ def _call():
             completion = self.client.chat.completions.create(
                 model=model_id or config["models"]["performance_low"],
                 temperature=temperature or config["temperature"]["low"],
-                max_tokens=max_tokens or config["llm"]["max_tokens"],
+                max_tokens=max_tokens or config["llm"]["max_tokens_output"],
                 response_format={
                     "type": "json_schema",
                     "json_schema": {
@@ -126,32 +130,71 @@ def _call():
 
         return _call()
 
-    def call_streamed(
+    def call_with_reasoning(
         self,
         prompt: str,
         model_id: str = None,
         temperature: float = None,
         max_tokens: int = None,
-        **kwargs,
-    ) -> Iterator:
-        """Stream responses from OpenRouter."""
+    ) -> tuple[str, dict]:
+        """Call LLM API with reasoning parameters."""
+
+        # At the moment, only the Gemini 2.5 model support context lengths beyond 200k.
+        # Here we check if the model is not Gemini 2.5 and if the token count exceeds the fallback limit.
+        # If so, we switch to the fallback model.
+        if "google/gemini-2.5" not in model_id or config["models"]["performance_high"]:
+            token_count = TokenCounter.count_tokens(prompt)
+            if token_count > config["llm"]["fallback_token_limit"]:
+                custom_logger.info_console(
+                    f"Token count ({token_count}) exceeds fallback limit. Using fallback model."
+                )
+                model_id = config["models"]["fallback"]
 
         custom_logger.info_console(
-            f"Streaming response for prompt: {prompt[:50]}... (model: {model_id or config['models']['performance_low']})"
+            f"Calling API with prompt: {prompt[:200]}... (model: {model_id or config['models']['performance_high']})"
         )
 
-        @self._retry
-        def _stream():
-            return self.client.chat.completions.create(
-                model=model_id or config["models"]["performance_low"],
-                temperature=temperature or config["temperature"]["low"],
-                max_tokens=max_tokens or config["llm"]["max_tokens"],
-                stream=True,
-                messages=[{"role": "user", "content": prompt}],
-                **kwargs,
-            )
-
-        return _stream()
+        # Since the final call is costly, we do not retry it, if it fails.
+        # If you want to retry it, uncomment the decorator below.
+        # @self._retry
+        def _call_model():
+            try:
+                url = "https://openrouter.ai/api/v1/chat/completions"
+                payload = {
+                    "model": model_id or config["models"]["performance_high"],
+                    "messages": [{"role": "user", "content": prompt}],
+                    "temperature": temperature or config["temperature"]["low"],
+                    "max_tokens": max_tokens or config["llm"]["max_tokens_output"],
+                    # Adjust this according to the model specifications. Details:
+                    # https://openrouter.ai/docs/use-cases/reasoning-tokens
+                    "reasoning": {
+                        "max_tokens": -1,
+                        # "effort": "high",
+                    },
+                }
+                headers = {
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/json",
+                }
+                response = requests.post(url, json=payload, headers=headers)
+
+                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+                with open(
+                    f"{config['app']['save_final_docs_to']}response_{timestamp}.json",
+                    "w",
+                ) as f:
+                    json.dump(response.json(), f)
+
+                response = response.json()
+                usage = response.get("usage", {})
+                response = response["choices"][0]["message"]["content"]
+                return response, usage
+            except Exception as e:
+                custom_logger.info_console(f"Error during final reasoning: {e}")
+                return "", {}
+
+        response, usage = _call_model()
+        return response, usage
 
 
 class ClientManager:

01_data/_weaviate_data/schema.db

@@ -1,7 +1,8 @@
 import pandas as pd
 import json
 import re
-from typing import List, Iterator, Union, Dict, Any
+from datetime import datetime
+from typing import List, Union, Dict, Any
 from _core.config import config
 from _core.models import ReflectTask, RelevanceCheck
 from _core.logger import custom_logger
@@ -224,8 +225,8 @@ def reflect_task_status(
 def create_final_report(
     user_query: str,
     final_docs: pd.DataFrame,
-    model_id: str = config["models"]["performance_medium"],
-) -> Iterator[str]:
+    model_id: str = config["models"]["performance_high"],
+) -> tuple[str, dict]:
     """Generate a final research report from selected documents."""
     research_results = [
         DOCUMENT.format(
@@ -241,14 +242,17 @@ def create_final_report(
     custom_logger.info_console(
         f"Creating final report for query: {user_query} with {len(research_results)} documents."
     )
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    with open(f"{config['app']['save_final_docs_to']}final_docs_{timestamp}.json", "w") as f:
+        json.dump(research_results, f, indent=2)
 
     research_results_text = "\n\n".join(research_results)
 
-    return llm_client.call_streamed(
+    response, usage = llm_client.call_with_reasoning(
         prompt=RESEARCH_WRITER.format(
             user_query=user_query, research_results=research_results_text
         ),
         model_id=model_id,
         temperature=config["temperature"]["base"],
-        reasoning_effort=config["llm"]["reasoning_effort"],
     )
+    return response, usage

02_app/_core/app_info.py

@@ -1,27 +1,3 @@
-INFO_TEXT = """Mit dieser App kannst du **vertieft über eigene Dokumentsammlungen recherchieren**.
-
-Die App dient zum Testen. **Beachte, dass Sprachmodelle (LLMs) Fehler machen und die Ergebnisse fehlerhaft oder unvollständig sein können.** Überprüfe die Ergebnisse immer.
-
-Deine Fragen werden an Clouddienste weitergeleitet und dort verarbeitet. **Gib daher nur als öffentlich klassifizierte Informationen als Fragen ein!** Beachte auch, dass die Nutzung anonymisiert aufgezeichnet wird und Mitarbeitende vom Statistischen Amt Eingaben stichprobenartig überprüfen, um die App zu verbessern.
-
-Zu Demonstrationszwecken bezieht die App für die Antworten eine kleine Auswahl von [Kantonsratsprotokollen des Kantons Zürich ein](https://opendata.swiss/de/dataset/zurcher-kantonsratsprotokolle-des-19-jahrhunderts). 
-
-Verantwortlich: Statistisches Amt, [Team Data](mailto:[email protected]).
-
-App-Version v0.1. Letzte Aktualisierung 17.7.2025
-
-### Wie funktioniert die App?
-
-Die App arbeitet in mehreren Schritten:
-
-1. **Suchanfragen formulieren**: Basierend auf deiner Fragestellung generiert die App gezielte Suchanfragen.  
-2. **Recherche durchführen**: Anschließend durchsucht die App die Dokumente nach passenden Textstellen.  
-3. **Relevanz prüfen**: Die gefundenen Passagen werden daraufhin geprüft, ob sie für deine Fragestellung von Bedeutung sind.  
-4. **Ganze Beschlüsse analysieren**: Zu relevanten Textstellen werden die gesamten Beschlüsse analysiert und inhaltlich zusammengefasst.  
-5. **Recherchestand reflektieren**: Die App bewertet den bisherigen Erkenntnisstand und entscheidet, ob eine weitere Recherche notwendig ist.  
-6. **Abschlussbericht erstellen**: Sobald ausreichend Informationen vorliegen oder die maximale Anzahl von drei Iterationen erreicht ist, fasst die App die Ergebnisse in einem Abschlussbericht zusammen.
-"""
-
 CREATE_QUERIES = """
 Du bist ein Rechercheassistent, spezialisiert auf Dokumente vom Kantonsrat Zürich.
 
@@ -86,10 +62,13 @@
 - Beziehe nur Informationen aus den Suchergebnissen ein, erfinde nichts.
 
 Hier ist die Frage bzw. Fragen des Experten:
+<expertenfrage>
 {user_query}
+</expertenfrage>
 
 Hier das Dokument vom Kanton Zürich:
 
+<beschluss>
 Titel
 {title}
 
@@ -101,20 +80,25 @@
 
 Dokument-Text
 {text}
+</beschluss>
 """.strip()
 
 
 DOCUMENT = """
+Regierungsratsbeschluss Kanton Zürich
 Dokument Kanton Zürich
 {title}
 {date}
 {link}
+
+Analyseergebnisse und Zusammenfassung der relevanten Informationen aus dem Dokument in Bezug auf die Frage(n) des Experten:
 {analysis}
 """.strip()
 
 
 REFLECT_TASK = """
 Du bist ein Rechercheassistent, spezialisiert auf Dokumente des Kantons Zürich.
+
 Deine Aufgabe ist es, den aktuellen Stand einer Recherche zu reflektieren und zu entscheiden, ob weitere Schritte erforderlich sind oder ob die Recherche abgeschlossen werden kann.
 
 Wichtige Hinweise:
@@ -129,15 +113,22 @@
     - False: Es sind weitere Schritte erforderlich, um die Fragen vollständig zu beantworten.
 
 Hier ist die Frage bzw. Fragen des Experten:
+<expertenfrage>
 {user_query}
+</expertenfrage>
 
 Hier sind die Analyseergebnisse von relevanten Dokumenten, die bisher erarbeitet wurden:
+<analyseergebnisse>
 {research_results}
-"""
+</analyseergebnisse>
+
+Reflektiere jetzt den aktuellen Stand der Recherche und entscheide, ob weitere Schritte erforderlich sind oder ob die Recherche abgeschlossen werden kann.
+""".strip()
 
 
 RESEARCH_WRITER = """
 Du bist ein Rechercheassistent, spezialisiert auf Dokumente vom Kanton Zürich.
+
 Deine Aufgabe ist es, die Ergebnisse einer Recherche in einem umfassenden, gut strukturierten Bericht zusammenzufassen.
 Du erhältst eine oder mehrere Fragen und eine Liste von Analyseergebnissen. Daraus sollst du einen Recherchebericht und präzise, juristisch fundierte Antworten erarbeiten.
 
@@ -172,8 +163,14 @@
     - Der Text endet mit dem Kapitel 3 Grundlagen und Quellen. Gib danach keine weiteren Kommentare oder Erklärungen ab.
 
 Hier ist die Frage bzw. die Fragen des Experten:
+<expertenfrage>
 {user_query}
+</expertenfrage>
 
 Hier sind die Analyseergebnisse von relevanten Dokumenten, die bisher erarbeitet wurden:
+<analyseergebnisse>
 {research_results}
-"""
+</analyseergebnisse>
+
+Erstelle jetzt den Recherchebericht und die Antworten auf die Fragen des Experten.
+""".strip()

02_app/_core/llm_client.py

@@ -82,7 +82,7 @@ def get_model_and_workflow_config(fast_mode=False):
             "check_relevance": config["models"]["performance_low"],
             "analyze_documents": config["models"]["performance_medium"],
             "reflect_task": config["models"]["performance_medium"],
-            "final_report": config["models"]["performance_medium"],
+            "final_report": config["models"]["performance_high"],
         }
         workflow_config = {
             "max_queries": config["app"]["max_queries"],