mirror of
https://github.com/lfnovo/open-notebook.git
synced 2026-04-28 11:30:00 +00:00
07c05ca354
Resolve conflicts in ask.py and chat.py by merging the try/except error handling from main with the extract_text_content helper from the PR. Also apply the same fix to source_chat.py and transformation.py which had the same vulnerable isinstance/str() pattern for structured LLM response content (e.g. Gemini's envelope format).
98 lines
3.5 KiB
Python
98 lines
3.5 KiB
Python
import asyncio
|
|
import sqlite3
|
|
from typing import Annotated, Optional
|
|
|
|
from ai_prompter import Prompter
|
|
from langchain_core.messages import AIMessage, 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)
|