chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,157 @@
|
||||
import operator
|
||||
from typing import Annotated, List
|
||||
|
||||
from ai_prompter import Prompter
|
||||
from langchain_core.output_parsers.pydantic import PydanticOutputParser
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.types import Send
|
||||
from pydantic import BaseModel, Field
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from open_notebook.ai.provision import provision_langchain_model
|
||||
from open_notebook.domain.notebook import vector_search
|
||||
from open_notebook.exceptions import OpenNotebookError
|
||||
from open_notebook.utils import clean_thinking_content
|
||||
from open_notebook.utils.error_classifier import classify_error
|
||||
from open_notebook.utils.text_utils import extract_text_content
|
||||
|
||||
|
||||
class SubGraphState(TypedDict):
|
||||
question: str
|
||||
term: str
|
||||
instructions: str
|
||||
results: dict
|
||||
answer: str
|
||||
ids: list # Added for provide_answer function
|
||||
|
||||
|
||||
class Search(BaseModel):
|
||||
term: str
|
||||
instructions: str = Field(
|
||||
description="Tell the answeting LLM what information you need extracted from this search"
|
||||
)
|
||||
|
||||
|
||||
class Strategy(BaseModel):
|
||||
reasoning: str
|
||||
searches: List[Search] = Field(
|
||||
default_factory=list,
|
||||
description="You can add up to five searches to this strategy",
|
||||
)
|
||||
|
||||
|
||||
class ThreadState(TypedDict):
|
||||
question: str
|
||||
strategy: Strategy
|
||||
answers: Annotated[list, operator.add]
|
||||
final_answer: str
|
||||
|
||||
|
||||
async def call_model_with_messages(state: ThreadState, config: RunnableConfig) -> dict:
|
||||
try:
|
||||
parser: PydanticOutputParser[Strategy] = PydanticOutputParser(
|
||||
pydantic_object=Strategy
|
||||
)
|
||||
system_prompt = Prompter(prompt_template="ask/entry", parser=parser).render( # type: ignore[arg-type]
|
||||
data=state # type: ignore[arg-type]
|
||||
)
|
||||
model = await provision_langchain_model(
|
||||
system_prompt,
|
||||
config.get("configurable", {}).get("strategy_model"),
|
||||
"tools",
|
||||
max_tokens=2000,
|
||||
structured=dict(type="json"),
|
||||
)
|
||||
# model = model.bind_tools(tools)
|
||||
# First get the raw response from the model
|
||||
ai_message = await model.ainvoke(system_prompt)
|
||||
|
||||
# Clean the thinking content from the response
|
||||
message_content = extract_text_content(ai_message.content)
|
||||
cleaned_content = clean_thinking_content(message_content)
|
||||
|
||||
# Parse the cleaned JSON content
|
||||
strategy = parser.parse(cleaned_content)
|
||||
|
||||
return {"strategy": strategy}
|
||||
except OpenNotebookError:
|
||||
raise
|
||||
except Exception as e:
|
||||
error_class, user_message = classify_error(e)
|
||||
raise error_class(user_message) from e
|
||||
|
||||
|
||||
async def trigger_queries(state: ThreadState, config: RunnableConfig):
|
||||
return [
|
||||
Send(
|
||||
"provide_answer",
|
||||
{
|
||||
"question": state["question"],
|
||||
"instructions": s.instructions,
|
||||
"term": s.term,
|
||||
# "type": s.type,
|
||||
},
|
||||
)
|
||||
for s in state["strategy"].searches
|
||||
]
|
||||
|
||||
|
||||
async def provide_answer(state: SubGraphState, config: RunnableConfig) -> dict:
|
||||
try:
|
||||
payload = state
|
||||
# if state["type"] == "text":
|
||||
# results = text_search(state["term"], 10, True, True)
|
||||
# else:
|
||||
results = await vector_search(state["term"], 10, True, True)
|
||||
if len(results) == 0:
|
||||
return {"answers": []}
|
||||
payload["results"] = results
|
||||
ids = [r["id"] for r in results]
|
||||
payload["ids"] = ids
|
||||
system_prompt = Prompter(prompt_template="ask/query_process").render(data=payload) # type: ignore[arg-type]
|
||||
model = await provision_langchain_model(
|
||||
system_prompt,
|
||||
config.get("configurable", {}).get("answer_model"),
|
||||
"tools",
|
||||
max_tokens=2000,
|
||||
)
|
||||
ai_message = await model.ainvoke(system_prompt)
|
||||
ai_content = extract_text_content(ai_message.content)
|
||||
return {"answers": [clean_thinking_content(ai_content)]}
|
||||
except OpenNotebookError:
|
||||
raise
|
||||
except Exception as e:
|
||||
error_class, user_message = classify_error(e)
|
||||
raise error_class(user_message) from e
|
||||
|
||||
|
||||
async def write_final_answer(state: ThreadState, config: RunnableConfig) -> dict:
|
||||
try:
|
||||
system_prompt = Prompter(prompt_template="ask/final_answer").render(data=state) # type: ignore[arg-type]
|
||||
model = await provision_langchain_model(
|
||||
system_prompt,
|
||||
config.get("configurable", {}).get("final_answer_model"),
|
||||
"tools",
|
||||
max_tokens=2000,
|
||||
)
|
||||
ai_message = await model.ainvoke(system_prompt)
|
||||
final_content = extract_text_content(ai_message.content)
|
||||
return {"final_answer": clean_thinking_content(final_content)}
|
||||
except OpenNotebookError:
|
||||
raise
|
||||
except Exception as e:
|
||||
error_class, user_message = classify_error(e)
|
||||
raise error_class(user_message) from e
|
||||
|
||||
|
||||
agent_state = StateGraph(ThreadState)
|
||||
agent_state.add_node("agent", call_model_with_messages)
|
||||
agent_state.add_node("provide_answer", provide_answer)
|
||||
agent_state.add_node("write_final_answer", write_final_answer)
|
||||
agent_state.add_edge(START, "agent")
|
||||
agent_state.add_conditional_edges("agent", trigger_queries, ["provide_answer"])
|
||||
agent_state.add_edge("provide_answer", "write_final_answer")
|
||||
agent_state.add_edge("write_final_answer", END)
|
||||
|
||||
graph = agent_state.compile()
|
||||
@@ -0,0 +1,98 @@
|
||||
import asyncio
|
||||
import sqlite3
|
||||
from typing import Annotated, Optional
|
||||
|
||||
from ai_prompter import Prompter
|
||||
from langchain_core.messages import SystemMessage
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.sqlite import SqliteSaver
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from open_notebook.ai.provision import provision_langchain_model
|
||||
from open_notebook.config import LANGGRAPH_CHECKPOINT_FILE
|
||||
from open_notebook.domain.notebook import Notebook
|
||||
from open_notebook.exceptions import OpenNotebookError
|
||||
from open_notebook.utils import clean_thinking_content
|
||||
from open_notebook.utils.error_classifier import classify_error
|
||||
from open_notebook.utils.text_utils import extract_text_content
|
||||
|
||||
|
||||
class ThreadState(TypedDict):
|
||||
messages: Annotated[list, add_messages]
|
||||
notebook: Optional[Notebook]
|
||||
context: Optional[str]
|
||||
context_config: Optional[dict]
|
||||
model_override: Optional[str]
|
||||
|
||||
|
||||
def call_model_with_messages(state: ThreadState, config: RunnableConfig) -> dict:
|
||||
try:
|
||||
system_prompt = Prompter(prompt_template="chat/system").render(data=state) # type: ignore[arg-type]
|
||||
payload = [SystemMessage(content=system_prompt)] + state.get("messages", [])
|
||||
model_id = config.get("configurable", {}).get("model_id") or state.get(
|
||||
"model_override"
|
||||
)
|
||||
|
||||
# Handle async model provisioning from sync context
|
||||
def run_in_new_loop():
|
||||
"""Run the async function in a new event loop"""
|
||||
new_loop = asyncio.new_event_loop()
|
||||
try:
|
||||
asyncio.set_event_loop(new_loop)
|
||||
return new_loop.run_until_complete(
|
||||
provision_langchain_model(
|
||||
str(payload), model_id, "chat", max_tokens=8192
|
||||
)
|
||||
)
|
||||
finally:
|
||||
new_loop.close()
|
||||
asyncio.set_event_loop(None)
|
||||
|
||||
try:
|
||||
# Try to get the current event loop
|
||||
asyncio.get_running_loop()
|
||||
# If we're in an event loop, run in a thread with a new loop
|
||||
import concurrent.futures
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||||
future = executor.submit(run_in_new_loop)
|
||||
model = future.result()
|
||||
except RuntimeError:
|
||||
# No event loop running, safe to use asyncio.run()
|
||||
model = asyncio.run(
|
||||
provision_langchain_model(
|
||||
str(payload),
|
||||
model_id,
|
||||
"chat",
|
||||
max_tokens=8192,
|
||||
)
|
||||
)
|
||||
|
||||
ai_message = model.invoke(payload)
|
||||
|
||||
# Clean thinking content from AI response (e.g., <think>...</think> tags)
|
||||
content = extract_text_content(ai_message.content)
|
||||
cleaned_content = clean_thinking_content(content)
|
||||
cleaned_message = ai_message.model_copy(update={"content": cleaned_content})
|
||||
|
||||
return {"messages": cleaned_message}
|
||||
except OpenNotebookError:
|
||||
raise
|
||||
except Exception as e:
|
||||
error_class, user_message = classify_error(e)
|
||||
raise error_class(user_message) from e
|
||||
|
||||
|
||||
conn = sqlite3.connect(
|
||||
LANGGRAPH_CHECKPOINT_FILE,
|
||||
check_same_thread=False,
|
||||
)
|
||||
memory = SqliteSaver(conn)
|
||||
|
||||
agent_state = StateGraph(ThreadState)
|
||||
agent_state.add_node("agent", call_model_with_messages)
|
||||
agent_state.add_edge(START, "agent")
|
||||
agent_state.add_edge("agent", END)
|
||||
graph = agent_state.compile(checkpointer=memory)
|
||||
@@ -0,0 +1,49 @@
|
||||
from typing import Any, Optional
|
||||
|
||||
from ai_prompter import Prompter
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from open_notebook.ai.provision import provision_langchain_model
|
||||
from open_notebook.utils.text_utils import clean_thinking_content, extract_text_content
|
||||
|
||||
|
||||
class PatternChainState(TypedDict):
|
||||
prompt: str
|
||||
parser: Optional[Any]
|
||||
input_text: str
|
||||
output: str
|
||||
|
||||
|
||||
async def call_model(state: dict, config: RunnableConfig) -> dict:
|
||||
content = state["input_text"]
|
||||
# state["prompt"] is caller-supplied free text. Never compile it as Jinja
|
||||
# template *source* (Prompter(template_text=...)) - pass it as a plain
|
||||
# render variable into a fixed, developer-authored template instead.
|
||||
# See docs/7-DEVELOPMENT/security.md (GHSA-f35w-wx37-26q7).
|
||||
system_prompt = Prompter(
|
||||
prompt_template="pattern/generic", parser=state.get("parser")
|
||||
).render(data=state)
|
||||
payload = [SystemMessage(content=system_prompt)] + [HumanMessage(content=content)]
|
||||
chain = await provision_langchain_model(
|
||||
str(payload),
|
||||
config.get("configurable", {}).get("model_id"),
|
||||
"transformation",
|
||||
max_tokens=5000,
|
||||
)
|
||||
|
||||
response = await chain.ainvoke(payload)
|
||||
|
||||
# Clean thinking tags from response (handles extended thinking models)
|
||||
output = clean_thinking_content(extract_text_content(response.content))
|
||||
return {"output": output}
|
||||
|
||||
|
||||
agent_state = StateGraph(PatternChainState)
|
||||
agent_state.add_node("agent", call_model) # type: ignore[type-var]
|
||||
agent_state.add_edge(START, "agent")
|
||||
agent_state.add_edge("agent", END)
|
||||
|
||||
graph = agent_state.compile()
|
||||
@@ -0,0 +1,246 @@
|
||||
import operator
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from content_core import ContentCoreConfig, extract_content
|
||||
from content_core.common import ExtractionOutput
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.types import Send
|
||||
from loguru import logger
|
||||
from typing_extensions import Annotated, TypedDict
|
||||
|
||||
from open_notebook.ai.models import Model, ModelManager
|
||||
from open_notebook.domain.content_settings import ContentSettings
|
||||
from open_notebook.domain.notebook import Asset, Source
|
||||
from open_notebook.domain.transformation import Transformation
|
||||
from open_notebook.graphs.transformation import graph as transform_graph
|
||||
|
||||
# Preferred languages for YouTube transcript selection. content-core's own
|
||||
# default is only ["en", "es", "pt"]; we keep the broader list Open Notebook has
|
||||
# always intended so non-English videos still resolve a transcript.
|
||||
YOUTUBE_PREFERRED_LANGUAGES = [
|
||||
"en",
|
||||
"pt",
|
||||
"es",
|
||||
"de",
|
||||
"nl",
|
||||
"en-GB",
|
||||
"fr",
|
||||
"hi",
|
||||
"ja",
|
||||
]
|
||||
|
||||
|
||||
class SourceState(TypedDict):
|
||||
# Input describing what to extract: url / file_path / content / delete_source.
|
||||
content_state: Dict[str, Any]
|
||||
# Result of content-core extraction (does NOT echo url/file_path back).
|
||||
extraction: ExtractionOutput
|
||||
apply_transformations: List[Transformation]
|
||||
source_id: str
|
||||
notebook_ids: List[str]
|
||||
source: Source
|
||||
transformation: Annotated[list, operator.add]
|
||||
embed: bool
|
||||
|
||||
|
||||
class TransformationState(TypedDict):
|
||||
source: Source
|
||||
transformation: Transformation
|
||||
|
||||
|
||||
async def content_process(state: SourceState) -> dict:
|
||||
content_state: Dict[str, Any] = state["content_state"]
|
||||
|
||||
# content-core 2.x takes engine/model overrides via ContentCoreConfig
|
||||
# (keyword-only), not inside the input dict.
|
||||
config_kwargs: Dict[str, Any] = {
|
||||
"youtube_languages": YOUTUBE_PREFERRED_LANGUAGES,
|
||||
}
|
||||
|
||||
# Honor the persisted content-processing engine choices. content-core
|
||||
# accepts "auto"/"simple"/"firecrawl"/"jina"/"crawl4ai" for URLs and
|
||||
# "auto"/"docling"/"simple" for documents; falling back to "auto" keeps the
|
||||
# previous behavior when settings are unset.
|
||||
try:
|
||||
settings: ContentSettings = await ContentSettings.get_instance() # type: ignore[assignment]
|
||||
if settings.default_content_processing_engine_url:
|
||||
config_kwargs["url_engine"] = settings.default_content_processing_engine_url
|
||||
if settings.default_content_processing_engine_doc:
|
||||
config_kwargs["document_engine"] = (
|
||||
settings.default_content_processing_engine_doc
|
||||
)
|
||||
if settings.docling_ocr is not None:
|
||||
config_kwargs["docling_ocr"] = settings.docling_ocr
|
||||
except Exception as e:
|
||||
# Keep the server-side traceback for diagnosing DB/deserialization
|
||||
# failures while still falling back to defaults (non-fatal).
|
||||
logger.opt(exception=True).warning(
|
||||
f"Failed to load content settings, using defaults: {e}"
|
||||
)
|
||||
|
||||
try:
|
||||
model_manager = ModelManager()
|
||||
defaults = await model_manager.get_defaults()
|
||||
if defaults.default_speech_to_text_model:
|
||||
stt_model = await Model.get(defaults.default_speech_to_text_model)
|
||||
if stt_model:
|
||||
config_kwargs["audio_provider"] = stt_model.provider
|
||||
config_kwargs["audio_model"] = stt_model.name
|
||||
logger.debug(
|
||||
f"Using speech-to-text model: {stt_model.provider}/{stt_model.name}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to retrieve speech-to-text model configuration: {e}")
|
||||
# Continue without custom audio model (content-core will use its default)
|
||||
|
||||
config = ContentCoreConfig(**config_kwargs) if config_kwargs else None
|
||||
|
||||
processed = await extract_content(
|
||||
url=content_state.get("url"),
|
||||
file_path=content_state.get("file_path"),
|
||||
content=content_state.get("content"),
|
||||
config=config,
|
||||
)
|
||||
|
||||
# content-core signals a soft extraction failure (e.g. an unreachable or
|
||||
# invalid URL, via the bs4 fallback) by returning title="Error" and content
|
||||
# prefixed with "Failed to extract content:" instead of raising. Detect that
|
||||
# sentinel and raise so the job is marked failed and the source becomes
|
||||
# retryable, rather than being saved as a "completed" source whose body is
|
||||
# the error string.
|
||||
if processed.title == "Error" and (processed.content or "").startswith(
|
||||
"Failed to extract content:"
|
||||
):
|
||||
raise ValueError(
|
||||
"Could not extract content from this source. "
|
||||
"The URL or file may be unreachable, invalid, or in an unsupported format."
|
||||
)
|
||||
|
||||
if not processed.content or not processed.content.strip():
|
||||
url = content_state.get("url") or ""
|
||||
if url and ("youtube.com" in url or "youtu.be" in url):
|
||||
raise ValueError(
|
||||
"Could not extract content from this YouTube video. "
|
||||
"No transcript or subtitles are available. "
|
||||
"Try configuring a Speech-to-Text model in Settings "
|
||||
"to transcribe the audio instead."
|
||||
)
|
||||
raise ValueError(
|
||||
"Could not extract any text content from this source. "
|
||||
"The content may be empty, inaccessible, or in an unsupported format."
|
||||
)
|
||||
|
||||
# content-core 2.x no longer deletes the uploaded source file after
|
||||
# extraction (the delete_source flag it used to honor is gone). Preserve the
|
||||
# previous auto-delete behavior on our side.
|
||||
if content_state.get("delete_source") and content_state.get("file_path"):
|
||||
file_path = content_state["file_path"]
|
||||
try:
|
||||
os.unlink(file_path)
|
||||
except FileNotFoundError:
|
||||
logger.warning(f"File not found while trying to delete: {file_path}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to delete source file {file_path}: {e}")
|
||||
|
||||
return {"extraction": processed}
|
||||
|
||||
|
||||
async def save_source(state: SourceState) -> dict:
|
||||
content_state = state["content_state"]
|
||||
extraction = state["extraction"]
|
||||
|
||||
# Get existing source using the provided source_id
|
||||
source = await Source.get(state["source_id"])
|
||||
if not source:
|
||||
raise ValueError(f"Source with ID {state['source_id']} not found")
|
||||
|
||||
# Update the source with processed content. content-core's ExtractionOutput
|
||||
# does not echo url/file_path back, so carry them from the input state.
|
||||
source.asset = Asset(
|
||||
url=content_state.get("url"), file_path=content_state.get("file_path")
|
||||
)
|
||||
source.full_text = extraction.content
|
||||
|
||||
# Preserve user-set title; only overwrite placeholder or empty titles
|
||||
if extraction.title and (not source.title or source.title == "Processing..."):
|
||||
source.title = extraction.title
|
||||
|
||||
await source.save()
|
||||
|
||||
# NOTE: Notebook associations are created by the API immediately for UI responsiveness
|
||||
# No need to create them here to avoid duplicate edges
|
||||
|
||||
if state["embed"]:
|
||||
if source.full_text and source.full_text.strip():
|
||||
logger.debug("Embedding content for vector search")
|
||||
await source.vectorize()
|
||||
else:
|
||||
logger.warning(
|
||||
f"Source {source.id} has no text content to embed, skipping vectorization"
|
||||
)
|
||||
|
||||
return {"source": source}
|
||||
|
||||
|
||||
def trigger_transformations(state: SourceState, config: RunnableConfig) -> List[Send]:
|
||||
if len(state["apply_transformations"]) == 0:
|
||||
return []
|
||||
|
||||
to_apply = state["apply_transformations"]
|
||||
logger.debug(f"Applying transformations {to_apply}")
|
||||
|
||||
return [
|
||||
Send(
|
||||
"transform_content",
|
||||
{
|
||||
"source": state["source"],
|
||||
"transformation": t,
|
||||
},
|
||||
)
|
||||
for t in to_apply
|
||||
]
|
||||
|
||||
|
||||
async def transform_content(state: TransformationState) -> Optional[dict]:
|
||||
source = state["source"]
|
||||
content = source.full_text
|
||||
if not content:
|
||||
return None
|
||||
transformation: Transformation = state["transformation"]
|
||||
|
||||
logger.debug(f"Applying transformation {transformation.name}")
|
||||
# LangGraph accepts a partial state dict at runtime, but its typed
|
||||
# overloads require the full state type (langgraph typing limitation).
|
||||
result = await transform_graph.ainvoke( # type: ignore[call-overload]
|
||||
dict(input_text=content, transformation=transformation)
|
||||
)
|
||||
await source.add_insight(transformation.title, result["output"])
|
||||
return {
|
||||
"transformation": [
|
||||
{
|
||||
"output": result["output"],
|
||||
"transformation_name": transformation.name,
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
# Create and compile the workflow
|
||||
workflow = StateGraph(SourceState)
|
||||
|
||||
# Add nodes
|
||||
workflow.add_node("content_process", content_process)
|
||||
workflow.add_node("save_source", save_source)
|
||||
workflow.add_node("transform_content", transform_content)
|
||||
# Define the graph edges
|
||||
workflow.add_edge(START, "content_process")
|
||||
workflow.add_edge("content_process", "save_source")
|
||||
workflow.add_conditional_edges(
|
||||
"save_source", trigger_transformations, ["transform_content"]
|
||||
)
|
||||
workflow.add_edge("transform_content", END)
|
||||
|
||||
# Compile the graph
|
||||
source_graph = workflow.compile()
|
||||
@@ -0,0 +1,254 @@
|
||||
import asyncio
|
||||
import sqlite3
|
||||
from typing import Annotated, Dict, List, Optional
|
||||
|
||||
from ai_prompter import Prompter
|
||||
from langchain_core.messages import SystemMessage
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.sqlite import SqliteSaver
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from open_notebook.ai.provision import provision_langchain_model
|
||||
from open_notebook.config import LANGGRAPH_CHECKPOINT_FILE
|
||||
from open_notebook.domain.notebook import Source, SourceInsight
|
||||
from open_notebook.exceptions import OpenNotebookError
|
||||
from open_notebook.utils import clean_thinking_content
|
||||
from open_notebook.utils.context_builder import build_source_context
|
||||
from open_notebook.utils.error_classifier import classify_error
|
||||
from open_notebook.utils.text_utils import extract_text_content
|
||||
|
||||
|
||||
class SourceChatState(TypedDict):
|
||||
messages: Annotated[list, add_messages]
|
||||
source_id: str
|
||||
source: Optional[Source]
|
||||
insights: Optional[List[SourceInsight]]
|
||||
context: Optional[str]
|
||||
model_override: Optional[str]
|
||||
context_indicators: Optional[Dict[str, List[str]]]
|
||||
|
||||
|
||||
def call_model_with_source_context(
|
||||
state: SourceChatState, config: RunnableConfig
|
||||
) -> dict:
|
||||
"""
|
||||
Main function that builds source context and calls the model.
|
||||
|
||||
This function:
|
||||
1. Uses build_source_context to build source-specific context
|
||||
2. Applies the source_chat Jinja2 prompt template
|
||||
3. Handles model provisioning with override support
|
||||
4. Tracks context indicators for referenced insights/content
|
||||
"""
|
||||
try:
|
||||
return _call_model_with_source_context_inner(state, config)
|
||||
except OpenNotebookError:
|
||||
raise
|
||||
except Exception as e:
|
||||
error_class, user_message = classify_error(e)
|
||||
raise error_class(user_message) from e
|
||||
|
||||
|
||||
def _call_model_with_source_context_inner(
|
||||
state: SourceChatState, config: RunnableConfig
|
||||
) -> dict:
|
||||
source_id = state.get("source_id")
|
||||
if not source_id:
|
||||
raise ValueError("source_id is required in state")
|
||||
|
||||
# Build source context using build_source_context (run async code in new loop)
|
||||
def build_context():
|
||||
"""Build context in a new event loop"""
|
||||
new_loop = asyncio.new_event_loop()
|
||||
try:
|
||||
asyncio.set_event_loop(new_loop)
|
||||
return new_loop.run_until_complete(
|
||||
build_source_context(
|
||||
source_id=source_id,
|
||||
max_tokens=50000, # Reasonable limit for source context
|
||||
)
|
||||
)
|
||||
finally:
|
||||
new_loop.close()
|
||||
asyncio.set_event_loop(None)
|
||||
|
||||
# Get the built context
|
||||
try:
|
||||
# Try to get the current event loop
|
||||
asyncio.get_running_loop()
|
||||
# If we're in an event loop, run in a thread with a new loop
|
||||
import concurrent.futures
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||||
future = executor.submit(build_context)
|
||||
context_data = future.result()
|
||||
except RuntimeError:
|
||||
# No event loop running, safe to create a new one
|
||||
context_data = build_context()
|
||||
|
||||
# Extract source and insights from context
|
||||
source = None
|
||||
insights = []
|
||||
context_indicators: dict[str, list[str | None]] = {
|
||||
"sources": [],
|
||||
"insights": [],
|
||||
"notes": [],
|
||||
}
|
||||
|
||||
if context_data.get("sources"):
|
||||
source_info = context_data["sources"][0] # First source
|
||||
source = Source(**source_info) if isinstance(source_info, dict) else source_info
|
||||
context_indicators["sources"].append(source.id)
|
||||
|
||||
if context_data.get("insights"):
|
||||
for insight_data in context_data["insights"]:
|
||||
insight = (
|
||||
SourceInsight(**insight_data)
|
||||
if isinstance(insight_data, dict)
|
||||
else insight_data
|
||||
)
|
||||
insights.append(insight)
|
||||
context_indicators["insights"].append(insight.id)
|
||||
|
||||
# Format context for the prompt
|
||||
formatted_context = _format_source_context(context_data)
|
||||
|
||||
# Build prompt data for the template
|
||||
prompt_data = {
|
||||
"source": source.model_dump() if source else None,
|
||||
"insights": [insight.model_dump() for insight in insights] if insights else [],
|
||||
"context": formatted_context,
|
||||
"context_indicators": context_indicators,
|
||||
}
|
||||
|
||||
# Apply the source_chat prompt template
|
||||
system_prompt = Prompter(prompt_template="source_chat/system").render(
|
||||
data=prompt_data
|
||||
)
|
||||
payload = [SystemMessage(content=system_prompt)] + state.get("messages", [])
|
||||
|
||||
# Handle async model provisioning from sync context
|
||||
def run_in_new_loop():
|
||||
"""Run the async function in a new event loop"""
|
||||
new_loop = asyncio.new_event_loop()
|
||||
try:
|
||||
asyncio.set_event_loop(new_loop)
|
||||
return new_loop.run_until_complete(
|
||||
provision_langchain_model(
|
||||
str(payload),
|
||||
config.get("configurable", {}).get("model_id")
|
||||
or state.get("model_override"),
|
||||
"chat",
|
||||
max_tokens=8192,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
new_loop.close()
|
||||
asyncio.set_event_loop(None)
|
||||
|
||||
try:
|
||||
# Try to get the current event loop
|
||||
asyncio.get_running_loop()
|
||||
# If we're in an event loop, run in a thread with a new loop
|
||||
import concurrent.futures
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||||
future = executor.submit(run_in_new_loop)
|
||||
model = future.result()
|
||||
except RuntimeError:
|
||||
# No event loop running, safe to use asyncio.run()
|
||||
model = asyncio.run(
|
||||
provision_langchain_model(
|
||||
str(payload),
|
||||
config.get("configurable", {}).get("model_id")
|
||||
or state.get("model_override"),
|
||||
"chat",
|
||||
max_tokens=8192,
|
||||
)
|
||||
)
|
||||
|
||||
ai_message = model.invoke(payload)
|
||||
|
||||
# Clean thinking content from AI response (e.g., <think>...</think> tags)
|
||||
content = extract_text_content(ai_message.content)
|
||||
cleaned_content = clean_thinking_content(content)
|
||||
cleaned_message = ai_message.model_copy(update={"content": cleaned_content})
|
||||
|
||||
# Update state with context information
|
||||
return {
|
||||
"messages": cleaned_message,
|
||||
"source": source,
|
||||
"insights": insights,
|
||||
"context": formatted_context,
|
||||
"context_indicators": context_indicators,
|
||||
}
|
||||
|
||||
|
||||
def _format_source_context(context_data: Dict) -> str:
|
||||
"""
|
||||
Format the context data into a readable string for the prompt.
|
||||
|
||||
Args:
|
||||
context_data: Context data from build_source_context
|
||||
|
||||
Returns:
|
||||
Formatted context string
|
||||
"""
|
||||
context_parts = []
|
||||
|
||||
# Add source information
|
||||
if context_data.get("sources"):
|
||||
context_parts.append("## SOURCE CONTENT")
|
||||
for source in context_data["sources"]:
|
||||
if isinstance(source, dict):
|
||||
context_parts.append(f"**Source ID:** {source.get('id', 'Unknown')}")
|
||||
context_parts.append(f"**Title:** {source.get('title', 'No title')}")
|
||||
if source.get("full_text"):
|
||||
# Truncate full text if too long
|
||||
full_text = source["full_text"]
|
||||
if len(full_text) > 5000:
|
||||
full_text = full_text[:5000] + "...\n[Content truncated]"
|
||||
context_parts.append(f"**Content:**\n{full_text}")
|
||||
context_parts.append("") # Empty line for separation
|
||||
|
||||
# Add insights
|
||||
if context_data.get("insights"):
|
||||
context_parts.append("## SOURCE INSIGHTS")
|
||||
for insight in context_data["insights"]:
|
||||
if isinstance(insight, dict):
|
||||
context_parts.append(f"**Insight ID:** {insight.get('id', 'Unknown')}")
|
||||
context_parts.append(
|
||||
f"**Type:** {insight.get('insight_type', 'Unknown')}"
|
||||
)
|
||||
context_parts.append(
|
||||
f"**Content:** {insight.get('content', 'No content')}"
|
||||
)
|
||||
context_parts.append("") # Empty line for separation
|
||||
|
||||
# Add metadata
|
||||
if context_data.get("metadata"):
|
||||
metadata = context_data["metadata"]
|
||||
context_parts.append("## CONTEXT METADATA")
|
||||
context_parts.append(f"- Source count: {metadata.get('source_count', 0)}")
|
||||
context_parts.append(f"- Insight count: {metadata.get('insight_count', 0)}")
|
||||
context_parts.append(f"- Total tokens: {context_data.get('total_tokens', 0)}")
|
||||
context_parts.append("")
|
||||
|
||||
return "\n".join(context_parts)
|
||||
|
||||
|
||||
# Create SQLite checkpointer
|
||||
conn = sqlite3.connect(
|
||||
LANGGRAPH_CHECKPOINT_FILE,
|
||||
check_same_thread=False,
|
||||
)
|
||||
memory = SqliteSaver(conn)
|
||||
|
||||
# Create the StateGraph
|
||||
source_chat_state = StateGraph(SourceChatState)
|
||||
source_chat_state.add_node("source_chat_agent", call_model_with_source_context)
|
||||
source_chat_state.add_edge(START, "source_chat_agent")
|
||||
source_chat_state.add_edge("source_chat_agent", END)
|
||||
source_chat_graph = source_chat_state.compile(checkpointer=memory)
|
||||
@@ -0,0 +1,13 @@
|
||||
from datetime import datetime
|
||||
|
||||
from langchain.tools import tool
|
||||
|
||||
|
||||
# todo: turn this into a system prompt variable
|
||||
@tool
|
||||
def get_current_timestamp() -> str:
|
||||
"""
|
||||
name: get_current_timestamp
|
||||
Returns the current timestamp in the format YYYYMMDDHHmmss.
|
||||
"""
|
||||
return datetime.now().strftime("%Y%m%d%H%M%S")
|
||||
@@ -0,0 +1,77 @@
|
||||
from ai_prompter import Prompter
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from open_notebook.ai.provision import provision_langchain_model
|
||||
from open_notebook.domain.notebook import Source
|
||||
from open_notebook.domain.transformation import DefaultPrompts, Transformation
|
||||
from open_notebook.exceptions import OpenNotebookError
|
||||
from open_notebook.utils import clean_thinking_content
|
||||
from open_notebook.utils.error_classifier import classify_error
|
||||
from open_notebook.utils.text_utils import extract_text_content
|
||||
|
||||
|
||||
class TransformationState(TypedDict):
|
||||
input_text: str
|
||||
source: Source
|
||||
transformation: Transformation
|
||||
output: str
|
||||
|
||||
|
||||
async def run_transformation(state: dict, config: RunnableConfig) -> dict:
|
||||
source_obj = state.get("source")
|
||||
source: Source = source_obj if isinstance(source_obj, Source) else None # type: ignore[assignment]
|
||||
content = state.get("input_text")
|
||||
assert source or content, "No content to transform"
|
||||
transformation: Transformation = state["transformation"]
|
||||
|
||||
try:
|
||||
if not content:
|
||||
content = source.full_text
|
||||
# transformation.prompt is user-controlled free text. Never compile it as
|
||||
# Jinja template *source* (Prompter(template_text=...)) - pass it as a
|
||||
# plain render variable into a fixed, developer-authored template instead.
|
||||
# See docs/7-DEVELOPMENT/security.md (GHSA-f35w-wx37-26q7).
|
||||
instructions = transformation.prompt
|
||||
default_prompts: DefaultPrompts = DefaultPrompts(transformation_instructions=None)
|
||||
if default_prompts.transformation_instructions:
|
||||
instructions = f"{default_prompts.transformation_instructions}\n\n{instructions}"
|
||||
|
||||
system_prompt = Prompter(prompt_template="transformation/execute").render(
|
||||
data={**state, "instructions": instructions}
|
||||
)
|
||||
content_str = str(content) if content else ""
|
||||
payload = [SystemMessage(content=system_prompt), HumanMessage(content=content_str)]
|
||||
chain = await provision_langchain_model(
|
||||
str(payload),
|
||||
config.get("configurable", {}).get("model_id"),
|
||||
"transformation",
|
||||
max_tokens=8192,
|
||||
)
|
||||
|
||||
response = await chain.ainvoke(payload)
|
||||
|
||||
# Clean thinking content from the response
|
||||
response_content = extract_text_content(response.content)
|
||||
cleaned_content = clean_thinking_content(response_content)
|
||||
|
||||
if source:
|
||||
await source.add_insight(transformation.title, cleaned_content)
|
||||
|
||||
return {
|
||||
"output": cleaned_content,
|
||||
}
|
||||
except OpenNotebookError:
|
||||
raise
|
||||
except Exception as e:
|
||||
error_class, user_message = classify_error(e)
|
||||
raise error_class(user_message) from e
|
||||
|
||||
|
||||
agent_state = StateGraph(TransformationState)
|
||||
agent_state.add_node("agent", run_transformation) # type: ignore[type-var]
|
||||
agent_state.add_edge(START, "agent")
|
||||
agent_state.add_edge("agent", END)
|
||||
graph = agent_state.compile()
|
||||
Reference in New Issue
Block a user