openai-chatbot-mcp / tools /vector_search_tool.py
Julian Vanecek
fixed ai stupid mistakes
285457c
"""
Vector Search Tool - Flexible search across one or more product versions.
"""
from typing import Dict, List, Optional
import json
def get_vector_search_tool_definition():
"""Get the OpenAI function definition for the vector search tool."""
return {
"type": "function",
"function": {
"name": "vector_search",
"description": "Search for information across one or more product versions. Use this to find specific information or compare across versions.",
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "The search query"
},
"versions": {
"type": "array",
"items": {
"type": "string",
"enum": ["harmony_1_2", "harmony_1_5", "harmony_1_6", "harmony_1_8", "chorus_1_1", "general_faq"]
},
"description": "List of versions to search. Can be a single version or multiple versions for comparison."
},
"max_results_per_version": {
"type": "integer",
"description": "Maximum results per version",
"default": 5,
"minimum": 1,
"maximum": 20
}
},
"required": ["query", "versions"]
}
}
}
def execute_vector_search(vector_store_manager, query: str, versions: List[str],
max_results_per_version: int = 5) -> Dict:
"""
Execute vector search across one or more versions.
Args:
vector_store_manager: Instance of VectorStoreManager
query: Search query
versions: List of version stores to search
max_results_per_version: Maximum results per version
Returns:
Dictionary with search results
"""
try:
# If single version, return simple format
if len(versions) == 1:
results = vector_store_manager.query_vector_store(
versions[0], query, max_results_per_version
)
if not results:
return {
"status": "success",
"message": f"No results found in {versions[0]}",
"results": []
}
return {
"status": "success",
"message": f"Found {len(results)} results in {versions[0]}",
"version": versions[0],
"query": query,
"results": results
}
# Multiple versions - return by version
all_results = {}
for version in versions:
results = vector_store_manager.query_vector_store(
version, query, max_results_per_version
)
if results:
all_results[version] = results
if not all_results:
return {
"status": "success",
"message": f"No results found across versions",
"results": {}
}
return {
"status": "success",
"message": f"Found results in {len(all_results)} versions",
"query": query,
"results": all_results
}
except Exception as e:
return {
"status": "error",
"message": f"Error searching: {str(e)}",
"results": {} if len(versions) > 1 else []
}
def format_search_results_for_context(results: Dict) -> str:
"""Format search results for including in the context."""
if results["status"] != "success":
return ""
# Handle single version search
if "version" in results and results.get("results"):
formatted = [f"Search results from {results['version']} for '{results['query']}':"]
for i, result in enumerate(results["results"], 1):
text = result.get("text", "")
similarity = result.get("similarity", 0)
formatted.append(f"{i}. {text} (similarity: {similarity:.3f})")
return "\n".join(formatted)
# Handle multi-version search
if isinstance(results.get("results"), dict):
formatted = [f"Search results across versions for '{results['query']}':"]
for version, version_results in results["results"].items():
formatted.append(f"\n{version}:")
for i, result in enumerate(version_results[:3], 1):
text = result.get("text", "")
similarity = result.get("similarity", 0)
formatted.append(f" {i}. {text} (similarity: {similarity:.3f})")
return "\n".join(formatted)
return ""
# Keep old function names for compatibility but have them use the new unified tool
def get_multi_version_search_tool_definition():
"""Deprecated - use get_vector_search_tool_definition instead."""
return get_vector_search_tool_definition()
def execute_multi_version_search(vector_store_manager, query: str, versions: List[str],
max_results_per_version: int = 3) -> Dict:
"""Deprecated - use execute_vector_search instead."""
return execute_vector_search(vector_store_manager, query, versions, max_results_per_version)