Fallbacks for failed structured output

langgraph.json

@@ -1,7 +1,7 @@
 {
     "dockerfile_lines": [],
     "graphs": {
-      "ollama_deep_researcher": "./src/assistant/graph.py:graph"
+      "ollama_deep_researcher": "./src/ollama_deep_researcher/graph.py:graph"
     },
     "python_version": "3.11",
     "env": "./.env",

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "ollama-deep-researcher"
 version = "0.0.1"
-description = "Lightweight web research and summarization assistant."
+description = "Fully local web research and summarization assistant with Ollama and LangGraph."
 authors = [
     { name = "Lance Martin" }
 ]
@@ -17,6 +17,7 @@ dependencies = [
     "beautifulsoup4>=4.13.3",
     "langchain-openai>=0.1.1",
     "openai>=1.12.0",
+    "langchain_openai>=0.3.9",
 ]
 
 [project.optional-dependencies]
@@ -27,10 +28,10 @@ requires = ["setuptools>=73.0.0", "wheel"]
 build-backend = "setuptools.build_meta"
 
 [tool.setuptools]
-packages = ["assistant"]
+packages = ["ollama_deep_researcher"]
 
 [tool.setuptools.package-dir]
-"assistant" = "src/assistant"
+"ollama_deep_researcher" = "src/ollama_deep_researcher"
 
 [tool.setuptools.package-data]
 "*" = ["py.typed"]

src/ollama_deep_researcher/__init__.py

@@ -0,0 +1 @@
+version = "0.0.1"

src/assistant/configuration.py → src/ollama_deep_researcher/configuration.py

@@ -1,20 +1,19 @@
 import os
-from typing import Any, Optional, Dict, List, Literal
+from enum import Enum
 from pydantic import BaseModel, Field
+from typing import Any, Optional, Literal
 
 from langchain_core.runnables import RunnableConfig
 
-from enum import Enum
-
 class SearchAPI(Enum):
     PERPLEXITY = "perplexity"
     TAVILY = "tavily"
     DUCKDUCKGO = "duckduckgo"
     SEARXNG = "searxng"
 
-@dataclass(kw_only=True)
-class Configuration:
+class Configuration(BaseModel):
     """The configurable fields for the research assistant."""
+
     max_web_research_loops: int = Field(
         default=3,
         title="Research Depth",
@@ -51,7 +50,7 @@ class Configuration:
         description="Base URL for LMStudio OpenAI-compatible API"
     )
     strip_thinking_tokens: bool = Field(
-        default=False,
+        default=True,
         title="Strip Thinking Tokens",
         description="Whether to strip <think> tokens from model responses"
     )

src/assistant/graph.py → src/ollama_deep_researcher/graph.py

@@ -7,12 +7,11 @@
 from langchain_ollama import ChatOllama
 from langgraph.graph import START, END, StateGraph
 
-from assistant.configuration import Configuration, SearchAPI
-from assistant.utils import deduplicate_and_format_sources, tavily_search, format_sources, perplexity_search, duckduckgo_search, searxng_search, strip_thinking_tokens
-
-from assistant.state import SummaryState, SummaryStateInput, SummaryStateOutput
-from assistant.prompts import query_writer_instructions, summarizer_instructions, reflection_instructions, get_current_date
-from assistant.lmstudio import ChatLMStudio
+from ollama_deep_researcher.configuration import Configuration, SearchAPI
+from ollama_deep_researcher.utils import deduplicate_and_format_sources, tavily_search, format_sources, perplexity_search, duckduckgo_search, searxng_search, strip_thinking_tokens, get_config_value
+from ollama_deep_researcher.state import SummaryState, SummaryStateInput, SummaryStateOutput
+from ollama_deep_researcher.prompts import query_writer_instructions, summarizer_instructions, reflection_instructions, get_current_date
+from ollama_deep_researcher.lmstudio import ChatLMStudio
 
 # Nodes
 def generate_query(state: SummaryState, config: RunnableConfig):
@@ -36,7 +35,7 @@ def generate_query(state: SummaryState, config: RunnableConfig):
             temperature=0, 
             format="json"
         )
-    else:  # Default to Ollama
+    else: # Default to Ollama
         llm_json_mode = ChatOllama(
             base_url=configurable.ollama_base_url, 
             model=configurable.local_llm, 
@@ -48,30 +47,29 @@ def generate_query(state: SummaryState, config: RunnableConfig):
         [SystemMessage(content=formatted_prompt),
         HumanMessage(content=f"Generate a query for web search:")]
     )
-    print(result.content)
 
-    # Strip thinking tokens if configured
+    # Get the content
     content = result.content
-    if configurable.strip_thinking_tokens:
-        content = strip_thinking_tokens(content)
-
-    query = json.loads(content)
 
-    return {"search_query": query['query']}
+    # Parse the JSON response and get the query
+    try:
+        query = json.loads(content)
+        search_query = query['query']
+    except (json.JSONDecodeError, KeyError):
+        # If parsing fails or the key is not found, use a fallback query
+        if configurable.strip_thinking_tokens:
+            content = strip_thinking_tokens(content)
+        search_query = content
+    return {"search_query": search_query}
 
 def web_research(state: SummaryState, config: RunnableConfig):
     """ Gather information from the web """
 
     # Configure
     configurable = Configuration.from_runnable_config(config)
 
-    # Handle both cases for search_api:
-    # 1. When selected in Studio UI -> returns a string (e.g. "tavily")
-    # 2. When using default -> returns an Enum (e.g. SearchAPI.TAVILY)
-    if isinstance(configurable.search_api, str):
-        search_api = configurable.search_api
-    else:
-        search_api = configurable.search_api.value
+    # Get the search API
+    search_api = get_config_value(configurable.search_api)
 
     # Search the web
     if search_api == "tavily":
@@ -135,14 +133,10 @@ def summarize_sources(state: SummaryState, config: RunnableConfig):
         HumanMessage(content=human_message_content)]
     )
 
+    # Strip thinking tokens if configured
     running_summary = result.content
-
-    # TODO: This is a hack to remove the <think> tags w/ Deepseek models
-    # It appears very challenging to prompt them out of the responses
-    while "<think>" in running_summary and "</think>" in running_summary:
-        start = running_summary.find("<think>")
-        end = running_summary.find("</think>") + len("</think>")
-        running_summary = running_summary[:start] + running_summary[end:]
+    if configurable.strip_thinking_tokens:
+        running_summary = strip_thinking_tokens(running_summary)
 
     return {"running_summary": running_summary}
 
@@ -160,7 +154,7 @@ def reflect_on_summary(state: SummaryState, config: RunnableConfig):
             temperature=0, 
             format="json"
         )
-    else:  # Default to Ollama
+    else: # Default to Ollama
         llm_json_mode = ChatOllama(
             base_url=configurable.ollama_base_url, 
             model=configurable.local_llm, 
@@ -174,36 +168,45 @@ def reflect_on_summary(state: SummaryState, config: RunnableConfig):
     )
 
     # Strip thinking tokens if configured
-    content = result.content
-    if configurable.strip_thinking_tokens:
-        content = strip_thinking_tokens(content)
-
-    follow_up_query = json.loads(content)
-
-    # Get the follow-up query
-    query = follow_up_query.get('follow_up_query')
-
-    # JSON mode can fail in some cases
-    if not query:
-        # Fallback to a placeholder query
+    try:
+        # Try to parse as JSON first
+        reflection_content = json.loads(result.content)
+        # Get the follow-up query
+        query = reflection_content.get('follow_up_query')
+        # Check if query is None or empty
+        if not query:
+            # Use a fallback query
+            return {"search_query": f"Tell me more about {state.research_topic}"}
+        return {"search_query": query}
+    except (json.JSONDecodeError, KeyError, AttributeError):
+        # If parsing fails or the key is not found, use a fallback query
         return {"search_query": f"Tell me more about {state.research_topic}"}
-
-    # Update search query with follow-up query
-    return {"search_query": follow_up_query['follow_up_query']}
-
+
 def finalize_summary(state: SummaryState):
     """ Finalize the summary """
 
-    # Format all accumulated sources into a single bulleted list
-    all_sources = "\n".join(source for source in state.sources_gathered)
+    # Deduplicate sources before joining
+    seen_sources = set()
+    unique_sources = []
+
+    for source in state.sources_gathered:
+        # Split the source into lines and process each individually
+        for line in source.split('\n'):
+            # Only process non-empty lines
+            if line.strip() and line not in seen_sources:
+                seen_sources.add(line)
+                unique_sources.append(line)
+
+    # Join the deduplicated sources
+    all_sources = "\n".join(unique_sources)
     state.running_summary = f"## Summary\n\n{state.running_summary}\n\n ### Sources:\n{all_sources}"
     return {"running_summary": state.running_summary}
 
 def route_research(state: SummaryState, config: RunnableConfig) -> Literal["finalize_summary", "web_research"]:
     """ Route the research based on the follow-up query """
 
     configurable = Configuration.from_runnable_config(config)
-    if state.research_loop_count <= int(configurable.max_web_research_loops):
+    if state.research_loop_count <= configurable.max_web_research_loops:
         return "web_research"
     else:
         return "finalize_summary"

src/assistant/lmstudio.py → src/ollama_deep_researcher/lmstudio.py

@@ -5,15 +5,10 @@
 from typing import Any, Dict, List, Optional, Union
 
 from langchain_core.callbacks.manager import CallbackManagerForLLMRun
-from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import (
-    AIMessage,
     BaseMessage,
-    ChatMessage,
-    HumanMessage,
-    SystemMessage,
 )
-from langchain_core.outputs import ChatGeneration, ChatResult
+from langchain_core.outputs import ChatResult
 from langchain_openai import ChatOpenAI
 from pydantic import Field
 

src/assistant/prompts.py → src/ollama_deep_researcher/prompts.py

@@ -18,15 +18,13 @@ def get_current_date():
 <FORMAT>
 Format your response as a JSON object with ALL three of these exact keys:
    - "query": The actual search query string
-   - "aspect": The specific aspect of the topic being researched
    - "rationale": Brief explanation of why this query is relevant
 </FORMAT>
 
 <EXAMPLE>
 Example output:
 {{
     "query": "machine learning transformer architecture explained",
-    "aspect": "technical architecture",
     "rationale": "Understanding the fundamental structure of transformer models"
 }}
 </EXAMPLE>

src/assistant/state.py → src/ollama_deep_researcher/state.py

@@ -1,6 +1,6 @@
 import operator
 from dataclasses import dataclass, field
-from typing_extensions import TypedDict, Annotated
+from typing_extensions import Annotated
 
 @dataclass(kw_only=True)
 class SummaryState:

src/assistant/utils.py → src/ollama_deep_researcher/utils.py

@@ -1,11 +1,29 @@
 import os
 import requests
-from typing import Dict, Any, List, Optional
+from typing import Dict, Any, List
 from langsmith import traceable
 from tavily import TavilyClient
 from duckduckgo_search import DDGS
 from langchain_community.utilities import SearxSearchWrapper
 
+def get_config_value(value):
+    """
+    Helper function to handle both string and enum cases of configuration values.
+
+    Args:
+        value: The configuration value to process. Can be a string or an Enum.
+
+    Returns:
+        str: The string representation of the value.
+
+    Examples:
+        >>> get_config_value("tavily")
+        'tavily'
+        >>> get_config_value(SearchAPI.TAVILY)
+        'tavily'
+    """
+    return value if isinstance(value, str) else value.value
+
 def strip_thinking_tokens(text: str) -> str:
     """Remove <think> tags and their content from the text.
 
@@ -15,6 +33,7 @@ def strip_thinking_tokens(text: str) -> str:
     Returns:
         str: The text with thinking tokens removed
     """
+
     while "<think>" in text and "</think>" in text:
         start = text.find("<think>")
         end = text.find("</think>") + len("</think>")
@@ -212,10 +231,7 @@ def tavily_search(query, include_raw_content=True, max_results=3):
                 - content (str): Snippet/summary of the content
                 - raw_content (str): Full content of the page if available"""
 
-    api_key = os.getenv("TAVILY_API_KEY")
-    if not api_key:
-        raise ValueError("TAVILY_API_KEY environment variable is not set")
-    tavily_client = TavilyClient(api_key=api_key)
+    tavily_client = TavilyClient()
     return tavily_client.search(query, 
                          max_results=max_results, 
                          include_raw_content=include_raw_content)