mirror of
https://github.com/imartinez/privateGPT.git
synced 2026-07-17 01:48:03 +00:00
091d5f7020
* fix: celery callbacks * fix: s3 + skill creator * fix: resumable when there's params * fix: add distributed cache * fix: do durable context stack * fix: mcp tools * fix: present server tool * fix: add cache to the skills * fix: mcp * fix: mypy
416 lines
16 KiB
Python
416 lines
16 KiB
Python
import asyncio
|
|
from typing import TYPE_CHECKING, Any, Literal, cast
|
|
|
|
from injector import inject, singleton
|
|
from llama_index.core.base.llms.types import (
|
|
ChatMessage,
|
|
MessageRole,
|
|
)
|
|
from llama_index.core.base.llms.types import (
|
|
TextBlock as LITextBlock,
|
|
)
|
|
|
|
from private_gpt.chat.input_models import BlobVisibilityMode
|
|
from private_gpt.components.cache import CacheService
|
|
from private_gpt.components.chat.models.chat_config_models import ToolSpec
|
|
from private_gpt.components.llm.llm_component import LLMComponent
|
|
from private_gpt.components.tools.binary_block_decorators import (
|
|
auto_resolve_media_blocks,
|
|
)
|
|
from private_gpt.components.tools.remote_execution import build_rebuild_metadata
|
|
from private_gpt.components.tools.tool_names import DATABASE_QUERY_TOOL_NAME
|
|
from private_gpt.components.tools.tool_placeholders import DATABASE_QUERY_TOOL_FN
|
|
from private_gpt.components.tools.types import ToolValidationMode
|
|
from private_gpt.di import get_global_injector
|
|
from private_gpt.events.models import (
|
|
BinaryBlock,
|
|
ResultContentBlockType,
|
|
TextBlock,
|
|
)
|
|
from private_gpt.server.utils.artifact_input import SqlDatabaseArtifact
|
|
from private_gpt.settings.settings import Settings
|
|
from private_gpt.utils.dependencies import format_missing_dependency_message
|
|
|
|
if TYPE_CHECKING:
|
|
from private_gpt.components.tabular.database_query_generator import (
|
|
DatabaseQueryGenerator,
|
|
ErrorQueryResult,
|
|
QueryResult,
|
|
)
|
|
|
|
|
|
def _load_database_query_dependencies() -> tuple[
|
|
type["DatabaseQueryGenerator"],
|
|
type["ErrorQueryResult"],
|
|
type["QueryResult"],
|
|
]:
|
|
try:
|
|
from private_gpt.components.tabular.database_query_generator import (
|
|
DatabaseQueryGenerator,
|
|
ErrorQueryResult,
|
|
QueryResult,
|
|
)
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
format_missing_dependency_message(
|
|
"Database query",
|
|
extras=(
|
|
"database-postgres",
|
|
"database-mysql",
|
|
"database-mssql",
|
|
"database-db2",
|
|
"database",
|
|
),
|
|
)
|
|
) from e
|
|
|
|
return DatabaseQueryGenerator, ErrorQueryResult, QueryResult
|
|
|
|
|
|
@singleton
|
|
class DatabaseQueryToolBuilder:
|
|
"""A builder class for creating a database query tool.
|
|
|
|
This tool allows users to run natural
|
|
language queries against connected SQL databases.
|
|
It runs the same query against all
|
|
provided databases and combines the results.
|
|
"""
|
|
|
|
@inject
|
|
def __init__(
|
|
self,
|
|
settings: Settings,
|
|
llm_component: LLMComponent,
|
|
cache: CacheService,
|
|
):
|
|
"""Initialize the DatabaseQueryToolBuilder with necessary components."""
|
|
self.settings = settings
|
|
self.llm_component = llm_component
|
|
self.cache = cache
|
|
self.sample_size = (
|
|
# number of characters to sample from the result for display
|
|
# TODO: this should be moved to tokens instead
|
|
500
|
|
)
|
|
|
|
@staticmethod
|
|
def _get_additional_context(
|
|
chat_history: list[ChatMessage] | None = None,
|
|
) -> str | None:
|
|
if not chat_history:
|
|
return None
|
|
|
|
additional_context: str = ""
|
|
|
|
# Add system prompt if available
|
|
system_prompt = next(
|
|
(msg for msg in chat_history if msg.role == MessageRole.SYSTEM),
|
|
None,
|
|
)
|
|
if system_prompt and system_prompt.content:
|
|
additional_context += (
|
|
f"**System Instructions**\n{system_prompt.content}\n\n"
|
|
)
|
|
|
|
# Add last user message if available
|
|
last_user_message = next(
|
|
(msg for msg in reversed(chat_history) if msg.role == MessageRole.USER),
|
|
None,
|
|
)
|
|
if last_user_message and last_user_message.content:
|
|
additional_context += (
|
|
f"**Current User Request\n{last_user_message.content}**\n\n"
|
|
)
|
|
|
|
# Add last database results if available
|
|
before_user_message = next(
|
|
(
|
|
msg
|
|
for msg in reversed(chat_history)
|
|
if msg.role == MessageRole.USER and (msg != last_user_message)
|
|
),
|
|
None,
|
|
)
|
|
last_iteration_messages = chat_history.copy() or []
|
|
if last_iteration_messages:
|
|
if last_user_message:
|
|
index = (
|
|
last_iteration_messages.index(last_user_message)
|
|
if last_user_message in last_iteration_messages
|
|
else -1
|
|
)
|
|
if index != -1:
|
|
last_iteration_messages = last_iteration_messages[:index]
|
|
|
|
if before_user_message:
|
|
index = (
|
|
last_iteration_messages.index(before_user_message)
|
|
if before_user_message in last_iteration_messages
|
|
else -1
|
|
)
|
|
if index != -1:
|
|
last_iteration_messages = last_iteration_messages[index:]
|
|
|
|
database_result_messages = [
|
|
msg
|
|
for msg in last_iteration_messages
|
|
if msg.role == MessageRole.TOOL
|
|
and msg.additional_kwargs.get("tool_call_name") == DATABASE_QUERY_TOOL_NAME
|
|
]
|
|
|
|
if database_result_messages:
|
|
# Take only the 1st block
|
|
# 1. Contains the query or an error
|
|
# 2. If apply, it contains the result summary
|
|
results = [
|
|
block.text
|
|
for message in database_result_messages
|
|
for block in message.blocks
|
|
if isinstance(block, LITextBlock)
|
|
]
|
|
if results:
|
|
additional_context += (
|
|
"** Previous Database Query (maybe it's relevant):**\n"
|
|
)
|
|
results_str = "\n".join(results)
|
|
additional_context += f"\n```\n{results_str}\n```\n\n"
|
|
|
|
additional_context = additional_context.strip()
|
|
return additional_context or None
|
|
|
|
async def build_tool(
|
|
self,
|
|
sql_artifacts: list[SqlDatabaseArtifact],
|
|
chat_history: list[ChatMessage] | None = None,
|
|
name: str = DATABASE_QUERY_TOOL_NAME,
|
|
type: str = DATABASE_QUERY_TOOL_NAME + "_v1",
|
|
description: str = DATABASE_QUERY_TOOL_FN.metadata.description,
|
|
validate: ToolValidationMode = ToolValidationMode.LAZY,
|
|
runtime: Literal["client", "server"] = "server",
|
|
blob_visibility: BlobVisibilityMode = BlobVisibilityMode.PUBLIC,
|
|
) -> ToolSpec:
|
|
(
|
|
database_query_generator_cls,
|
|
error_query_result_cls,
|
|
query_result_cls,
|
|
) = _load_database_query_dependencies()
|
|
|
|
sample_size = self.sample_size # capture for closure
|
|
|
|
async def validate_sql() -> None:
|
|
if not sql_artifacts:
|
|
raise ValueError("At least one SQL database artifact is required.")
|
|
|
|
query_gen = [
|
|
database_query_generator_cls(
|
|
connection_string=sql_artifact.connection_string,
|
|
ssl=sql_artifact.ssl,
|
|
schemas=sql_artifact.schemas,
|
|
# TODO: Meanwhile we validate the feature,
|
|
# we set readonly to true to avoid accidental data changes
|
|
is_readonly=True,
|
|
enable_tables=sql_artifact.enable_tables,
|
|
enable_views=sql_artifact.enable_views,
|
|
enable_functions=sql_artifact.enable_functions,
|
|
enable_procedures=sql_artifact.enable_procedures,
|
|
description=sql_artifact.description,
|
|
batch_size=self.settings.database_query.batch_size,
|
|
timeout_seconds=self.settings.database_query.timeout_seconds,
|
|
max_mb_result=self.settings.database_query.max_mb_result,
|
|
cache=self.cache,
|
|
)
|
|
for sql_artifact in sql_artifacts
|
|
]
|
|
validations = [gen.check_connection() for gen in query_gen]
|
|
results = list(await asyncio.gather(*validations, return_exceptions=True))
|
|
errors = [
|
|
f"- DB {i} connection error: {result!s}\n"
|
|
for i, result in enumerate(results)
|
|
if isinstance(result, str | Exception)
|
|
]
|
|
if errors:
|
|
raise ConnectionError(
|
|
"One or more database connections failed: \n" + "; ".join(errors),
|
|
)
|
|
|
|
@auto_resolve_media_blocks(blob_visibility=blob_visibility)
|
|
async def execute_sql(query: str) -> list[ResultContentBlockType]:
|
|
additional_context: str | None = await asyncio.to_thread(
|
|
self._get_additional_context,
|
|
chat_history,
|
|
)
|
|
query_gen = []
|
|
try:
|
|
query_gen = [
|
|
database_query_generator_cls(
|
|
connection_string=sql_artifact.connection_string,
|
|
ssl=sql_artifact.ssl,
|
|
schemas=sql_artifact.schemas,
|
|
enable_tables=sql_artifact.enable_tables,
|
|
enable_views=sql_artifact.enable_views,
|
|
enable_functions=sql_artifact.enable_functions,
|
|
enable_procedures=sql_artifact.enable_procedures,
|
|
description=sql_artifact.description,
|
|
batch_size=self.settings.database_query.batch_size,
|
|
timeout_seconds=self.settings.database_query.timeout_seconds,
|
|
max_mb_result=self.settings.database_query.max_mb_result,
|
|
cache=self.cache,
|
|
)
|
|
for sql_artifact in sql_artifacts
|
|
]
|
|
|
|
# create the query coroutines
|
|
searches = [
|
|
gen.query(
|
|
query=query,
|
|
additional_context=additional_context,
|
|
)
|
|
for gen in query_gen
|
|
]
|
|
query_results_or_exceptions = await asyncio.gather(
|
|
*searches, return_exceptions=True
|
|
)
|
|
query_results = []
|
|
for result in query_results_or_exceptions:
|
|
if isinstance(result, BaseException):
|
|
query_results.append(
|
|
query_result_cls(
|
|
query=None,
|
|
row_count=-1,
|
|
error=error_query_result_cls(str(result)),
|
|
)
|
|
)
|
|
else:
|
|
query_results.append(result)
|
|
if not query_results:
|
|
return [TextBlock(text="No databases returned results.")]
|
|
|
|
# if all row counts are 0, return no results found
|
|
all_failed = all(result.row_count < 0 for result in query_results)
|
|
if all_failed:
|
|
return [
|
|
TextBlock(text=f"No results found.\nError: {result.error}")
|
|
for result in query_results
|
|
]
|
|
|
|
results = list(zip(sql_artifacts, query_results, strict=False))
|
|
query_with_results = [
|
|
(sql_artifact, db_query_result)
|
|
for sql_artifact, db_query_result in results
|
|
if db_query_result.row_count > 0
|
|
]
|
|
if query_with_results:
|
|
# As the model cannot understand that some
|
|
# databases returned no results,
|
|
# we only show the databases that returned results
|
|
# if at least one database returned results
|
|
results = query_with_results
|
|
|
|
result_as_block_list: list[list[ResultContentBlockType]] = []
|
|
for sql_artifact, db_query_result in results:
|
|
prefix = (
|
|
f"Database: {sql_artifact.connection_string}\n"
|
|
if len(results) > 1
|
|
else ""
|
|
)
|
|
blocks: list[ResultContentBlockType] = [
|
|
TextBlock(
|
|
text=prefix
|
|
+ "Query:\n```sql\n"
|
|
+ (db_query_result.query or "No query was generated.")
|
|
+ "\n```\n\n Row Count: "
|
|
+ str(db_query_result.row_count)
|
|
)
|
|
]
|
|
|
|
if db_query_result.warning:
|
|
blocks.append(
|
|
TextBlock(text="Warning: \n" + str(db_query_result.warning))
|
|
)
|
|
|
|
if db_query_result.row_count > 0:
|
|
csv = db_query_result.as_csv()
|
|
filename = f"csv_{hash(query)}.csv"
|
|
csv_block = BinaryBlock.from_text(
|
|
text=csv,
|
|
filename=filename,
|
|
mime_type="text/csv",
|
|
)
|
|
blocks.append(csv_block)
|
|
|
|
if len(csv) > sample_size:
|
|
blocks.append(
|
|
TextBlock(
|
|
text="Representative data from the query result. Information IS NOT COMPLETE, refer"
|
|
"to the generated csv file for the full response:\n"
|
|
+ (csv[0:sample_size])
|
|
)
|
|
)
|
|
else:
|
|
blocks.append(TextBlock(text="Query result: \n" + csv))
|
|
elif db_query_result.query and db_query_result.row_count == 0:
|
|
blocks.append(
|
|
TextBlock(
|
|
text="Query executed successfully. No rows found."
|
|
)
|
|
)
|
|
elif db_query_result.error:
|
|
blocks.append(
|
|
TextBlock(text="Error: \n" + str(db_query_result.error))
|
|
)
|
|
|
|
result_as_block_list.append(blocks)
|
|
|
|
flattened_blocks: list[ResultContentBlockType] = [
|
|
block for sublist in result_as_block_list for block in sublist
|
|
]
|
|
if not flattened_blocks:
|
|
flattened_blocks = [
|
|
TextBlock(text="No results found for the query.")
|
|
]
|
|
|
|
return flattened_blocks
|
|
finally:
|
|
# close connections
|
|
def close() -> None:
|
|
for gen in query_gen:
|
|
gen.close()
|
|
|
|
await asyncio.to_thread(close)
|
|
|
|
async def run_tool(query: str) -> list[ResultContentBlockType]:
|
|
if validate == ToolValidationMode.LAZY:
|
|
await validate_sql()
|
|
|
|
return await execute_sql(query)
|
|
|
|
if validate == ToolValidationMode.EAGER:
|
|
await validate_sql()
|
|
|
|
return ToolSpec.from_defaults(
|
|
name=name,
|
|
type=type,
|
|
runtime=runtime,
|
|
description=description,
|
|
async_fn=run_tool,
|
|
execution_metadata=build_rebuild_metadata(
|
|
rebuild_database_query_tool,
|
|
{
|
|
"sql_artifacts": sql_artifacts,
|
|
"chat_history": chat_history,
|
|
"name": name,
|
|
"type": type,
|
|
"description": description,
|
|
"validate": validate,
|
|
"runtime": runtime,
|
|
"blob_visibility": blob_visibility,
|
|
},
|
|
),
|
|
)
|
|
|
|
|
|
async def rebuild_database_query_tool(**kwargs: Any) -> ToolSpec:
|
|
builder = get_global_injector().get(DatabaseQueryToolBuilder)
|
|
return await builder.build_tool(**cast(Any, kwargs))
|