Core Module

Main Module

Agent Graph

State Management

State classes for OpenChatBI graph execution.

openchatbi.graph_state.add_history_messages(left: list, right: list)[source]

class openchatbi.graph_state.AgentState[source]

Bases: MessagesState

State for the main agent graph execution.

Extends MessagesState with additional fields for routing and responses.

history_messages: AIMessage], <function add_history_messages at 0x7fe144de2980>]

agent_next_node: str

sends: list[Send]

sql: str

final_answer: str

messages: _add_messages at 0x7fe154fac5e0>]

class openchatbi.graph_state.SQLGraphState[source]

Bases: MessagesState

State for SQL generation subgraph.

Contains rewritten question, table selection, extracted entities, and generated SQL.

rewrite_question: str

tables: list[dict[str, Any]]

info_entities: dict[str, Any]

sql: str

sql_retry_count: int

sql_execution_result: str

schema_info: dict[str, Any]

data: str

previous_sql_errors: list[dict[str, Any]]

visualization_dsl: dict[str, Any]

messages: _add_messages at 0x7fe154fac5e0>]

class openchatbi.graph_state.InputState[source]

Bases: MessagesState

Input state schema for the main graph.

messages: _add_messages at 0x7fe154fac5e0>]

class openchatbi.graph_state.OutputState[source]

Bases: MessagesState

Output state schema for the main graph.

messages: _add_messages at 0x7fe154fac5e0>]

class openchatbi.graph_state.SQLOutputState[source]

Bases: MessagesState

Output state schema for the SQL generation subgraph.

rewrite_question: str

tables: list[dict[str, Any]]

sql: str

schema_info: dict[str, Any]

data: str

visualization_dsl: dict[str, Any]

messages: _add_messages at 0x7fe154fac5e0>]

Utilities

Utility functions for OpenChatBI.

openchatbi.utils.log(args) → None[source]: Log messages to stderr for debugging.

openchatbi.utils.get_text_from_content(content: str | list[str | dict]) → str[source]

Extract text from various content formats.

Parameters:: content – String, list of strings, or list of dicts with ‘text’ key.
Returns:: Extracted text content.
Return type:: str

openchatbi.utils.get_text_from_message_chunk(chunk: AIMessageChunk) → str[source]

Extract content from an AIMessageChunk.

Parameters:: chunk (AIMessageChunk) – The message chunk to extract text from.
Returns:: Extracted text content or empty string.
Return type:: str

openchatbi.utils.extract_json_from_answer(answer: str) → dict[source]

Extract the first JSON object from a string answer.

Parameters:: answer (str) – String that may contain JSON objects.
Returns:: Parsed JSON object or empty dict if none found.
Return type:: dict

openchatbi.utils.get_report_download_response(filename: str) → FileResponse[source]

Get FileResponse for downloading a report file.

Parameters:: filename – The filename of the report to download
Returns:: Response object for file download
Return type:: FileResponse
Raises:: HTTPException – Various HTTP errors for invalid requests

openchatbi.utils.create_vector_db(texts: list[str], embedding=None, collection_name: str = 'langchain', metadatas=None, collection_metadata: dict = None) → VectorStore[source]

Create or reuse a Chroma vector database.

Parameters:

texts (List[str]) – Text documents to index.
embedding – Embedding function to use.
collection_name (str) – Name of the collection.
metadatas – Metadata for each document.
collection_metadata (dict) – Collection-level metadata.

Returns:

Vector database instance.

Return type:

Chroma

openchatbi.utils.recover_incomplete_tool_calls(state: AgentState) → list[source]

Recover from incomplete tool calls by creating message operations to insert ToolMessages correctly.

When the graph execution is interrupted (e.g., by kill or app restart) during tool execution, the state can end up with AIMessage containing tool_calls but no corresponding ToolMessage responses. This function detects such cases and creates the necessary message operations to insert failure ToolMessages in the correct position (right after the AIMessage).

Parameters:: state (AgentState) – The current graph state containing messages.
Returns:: Message operations to insert recovery ToolMessages, or empty list if no recovery needed.
Return type:: list