digichat / agent_config.py
chrizefan's picture
Upload folder using huggingface_hub
fe52ef9 verified
Raw
History Blame Contribute Delete
31.5 kB
from typing import Dict, Any, List
import os
import json
import datetime
from abc import ABC, abstractmethod
class AgentConfig(ABC):
"""Base configuration class for agents"""
def __init__(self):
self.data_sources: Dict[str, Dict[str, Any]] = {}
self.tools: List[Dict[str, Any]] = []
self.system_prompts: Dict[str, str] = {}
self.model_config: Dict[str, Any] = {
"model": "grok-2-latest",
"temperature": 0.5,
"max_iterations": 10
}
self.azure_index: Dict[str, Any] = {}
self.tool_agent_map: Dict[str, str] = {} # tool name -> agent type
@abstractmethod
def setup_data_sources(self) -> None:
"""Setup data source configurations"""
pass
@abstractmethod
def setup_tools(self) -> None:
"""Setup available tools"""
pass
@abstractmethod
def setup_system_prompts(self) -> None:
"""Setup system prompts"""
pass
@abstractmethod
def setup_azure_index(self) -> None:
"""Setup search indexes configurations"""
pass
def setup_model_config(self, model: str = "grok-2-latest", temperature: float = 0.5, max_iterations: int = 10) -> None:
"""Setup model configuration"""
self.model_config = {
"model": model,
"temperature": temperature,
"max_iterations": max_iterations
}
def get_data_sources(self) -> Dict[str, Dict[str, Any]]:
"""Get all data source configurations"""
return self.data_sources
def get_azure_index(self) -> Dict[str, Any]:
"""Get all search index configurations"""
return self.azure_index
def get_tools(self) -> List[Dict[str, Any]]:
"""Get the list of available tools"""
return self.tools
def get_system_prompt(self, agent_type: str) -> str:
"""Get the system prompt for a specific agent type"""
if agent_type not in self.system_prompts:
return None
return self.system_prompts[agent_type]
def get_model_config(self) -> Dict[str, Any]:
"""Get the model configuration"""
return self.model_config
def to_json(self) -> str:
"""Convert the configuration to a JSON string"""
config = {
"tools": self.tools,
"system_prompts": self.system_prompts,
"model_config": self.model_config,
"data_sources": self.data_sources,
"azure_index": self.azure_index
}
return json.dumps(config, indent=2)
def get_azure_index(self) -> Dict[str, Any]:
"""Get the Azure Search configuration"""
return self.azure_index
def get_tool_agent_map(self) -> Dict[str, str]:
"""Get the mapping from tool name to agent type"""
return self.tool_agent_map
class BaseAgentConfig(AgentConfig):
"""Standard configuration class for basic language model agents"""
def __init__(self, model: str = "grok-2-latest", temperature: float = 0.5):
super().__init__()
self.model = model
self.setup_model_config(model, temperature)
self.setup_data_sources()
self.setup_tools()
self.setup_system_prompts()
self.setup_azure_index()
def setup_data_sources(self) -> None:
"""Basic agent has no data sources"""
self.data_sources = {}
def setup_tools(self) -> None:
"""Basic agent has no tools"""
self.tools = []
def setup_system_prompts(self) -> None:
"""Setup standard system prompt"""
self.system_prompts = {
"default": "You are a helpful AI assistant."
}
def setup_azure_index(self) -> None:
"""Basic agent has no search indexes"""
self.azure_index = {}
class SITAASAgentConfig(BaseAgentConfig):
"""Configuration class for the SITAAS agent"""
def __init__(self, model: str = "gpt-4o-mini", temperature: float = 0.0):
super().__init__(model, temperature)
self.setup_tool_agent_map()
def setup_data_sources(self) -> None:
"""Setup SITAAS data source configurations"""
self.data_sources = {
"preview_metadata_agent": {
"type": "azure_search",
"parameters": {
"endpoint": os.getenv("SEARCH_ENDPOINT"),
"index_name": "azureblob-index",
"semantic_configuration": "default",
"query_type": "simple",
"fields_mapping": {},
"in_scope": True,
"filter": None,
"strictness": 1,
"top_n_documents": 20,
"authentication": {
"type": "api_key",
"key": os.getenv("SEARCH_KEY"),
}
}
},
"preview_content_agent": {
"type": "azure_search",
"parameters": {
"endpoint": os.getenv("SEARCH_ENDPOINT"),
"index_name": "vector-1745685878660",
"semantic_configuration": "default",
"query_type": "vector_simple_hybrid",
"fields_mapping": {},
"in_scope": True,
"filter": None,
"strictness": 1,
"top_n_documents": 20,
"authentication": {
"type": "api_key",
"key": os.getenv("SEARCH_KEY"),
},
"embedding_dependency": {
"type": "deployment_name",
"deployment_name": "text-embedding-ada-002"
}
}
}
}
def setup_tools(self) -> None:
"""Setup SITAAS available tools"""
self.tools = [
{
"type": "function",
"function": {
"name": "search_metadata",
"description": "Retrieve metadata documents from Azure Cognitive Search using Lucene syntax. Useful when more filtering or exploration of metadata is needed.",
"parameters": {
"type": "object",
"properties": {
"search_text": {
"type": "string",
"description": "Lucene syntax query to retrieve metadata for documents.",
"default": "*",
},
"num_results": {
"type": "integer",
"description": "Number of metadata records to retrieve.",
"minimum": 1,
"maximum": 1000
},
"filter": {
"type": "string",
"description": "Optional OData filter expression to apply to the search.",
},
"orderby": {
"type": "string",
"description": "Optional OData orderby expression to sort the results.",
},
"select": {
"type": "string",
"description": "Optional OData select expression to limit the fields returned in the results.",
"default": "Id,Owner,Name,WebUrl,CreatedDataTime,CreatedBy,LastModifiedDate,LastModifiedBy"
},
},
"required": ["search_text"],
"additionalProperties": False
}
}
},
{
"type": "function",
"function": {
"name": "search_content",
"description": "Similarity search on document chunks in a vector index. Ideal for high-recall document content analysis.",
"parameters": {
"type": "object",
"properties": {
"search_text": {
"type": "string",
"description": "Semantic query for content-based document retrieval.",
"default": "*",
},
"num_results": {
"type": "integer",
"description": "Number of document chunks to retrieve.",
"minimum": 1,
"maximum": 1000
},
"filter": {
"type": "string",
"description": "Optional OData filter expression to apply to the search.",
},
"orderby": {
"type": "string",
"description": "Optional OData orderby expression to sort the results.",
},
"select": {
"type": "string",
"description": "Optional OData select expression to limit the fields returned in the results.",
"default": "title,chunk_id,chunk, Id, Owner, Name"
},
},
"required": ["search_text"],
"additionalProperties": False
}
}
},
# {
# "type": "function",
# "function": {
# "name": "generate_final_response",
# "description": "Generate a final synthesized answer using the selected documents and an optional formatting instruction.",
# "parameters": {
# "type": "object",
# "properties": {
# "format_instruction": {
# "type": "string",
# "description": "Optional formatting guide for the final output (e.g., summary, markdown, table, Q&A, etc.)."
# }
# },
# "additionalProperties": False
# }
# }
# }
]
def setup_azure_index(self) -> None:
"""Setup SITAAS search indexes configurations"""
self.azure_index = {
"search_endpoint": os.getenv("SEARCH_ENDPOINT"),
"search_key": os.getenv("SEARCH_KEY"),
"search_index_name": "azureblob-index",
"vector_index_name": "vector-1745685878660",
"embeddings_deployment": "text-embedding-ada-002",
"azure_openai_endpoint": os.getenv("AZURE_OPENAI_ENDPOINT"),
"azure_openai_key": os.getenv("AZURE_OPENAI_API_KEY")
}
def setup_system_prompts(self) -> None:
"""Setup SITAAS system prompts"""
self.system_prompts = {
"orchestrator": f"""
You are an intelligent orchestration agent designed to fulfill complex user queries by coordinating across multiple specialized document search tools and output generation utilities.
You are plugged into a database and have access to powerful search tools. If the user's request is ambiguous or missing critical information, you may ask a clarifying question, but do not overdo it—prefer to take action and use your tools to retrieve information whenever possible.
Only ask the user for more details if it is truly necessary to proceed. Otherwise, attempt to fulfill the request using the available tools and data.
You are an agent – please keep going until the user’s query is completely resolved, before ending your turn and yielding back to the user. Only terminate your turn when you are sure that the problem is solved.
If you are not sure about file content or document structure pertaining to the user’s request, use your tools to retrieve information. Do NOT guess or make up an answer.
You MUST plan before each function call and reflect on the outcomes of previous function calls, but do not delay action by repeatedly asking for clarification.
**Date**: Today is {datetime.datetime.now().strftime("%Y-%m-%d")}.
---
### 🧠 Planning & Reflection
Before each tool call, briefly **plan** out your next action.
After each tool call, **reflect** on what was retrieved and adjust your next step accordingly.
Do not simply chain tool calls without commentary or thinking.
### 🔧 Tool-Calling Reminder
Use your tools to gather information instead of guessing.
Each tool serves a purpose—select them wisely based on the query type.
### 🔍 Search Strategy
1. **Scoping Phase**:
- Start by calling search tools (such as preview or metadata tools) with a small number of results (e.g., 5–10).
- Use these initial results to understand the structure, quality, and relevance of the data.
- Analyze the returned results to identify useful filters, patterns, or query refinements.
- If results are ambiguous or insufficient, adjust your search parameters (e.g., keywords, filters, or sort order) and try again with a small result set.
2. **Refinement Phase**:
- Based on insights from the scoping phase, refine your search parameters to target the most relevant data.
- Continue using small result sets until you are confident that your search parameters are well-tuned and will yield high-quality results.
3. **Comprehensive Search Phase**:
- Once you have refined your search parameters and are confident in their effectiveness, perform a comprehensive search without a strict result limit (or with a much higher limit).
- Retrieve all relevant documents or content needed to fully address the user's query.
4. **Final Output**:
- Once you’ve gathered sufficient content, use the `generate_final_response` tool to produce a human-readable output.
- If the user specified a desired output format (e.g., "summary", "bullet points", "explanatory"), pass that to the response tool.
---
### 🧩 Metadata Fields Available
Use these fields for Lucene-based searches and filters on metadata:
| **Field** | **Description** | 🔍 **Searchable** | 🔎 **Filterable** |
|----------------------------------|--------------------------------------------------------------------------------------------------|------------------|--------------------|
| `Id` | Unique identifier for the metadata record. | ✅ | ✅ |
| `Name` | The name of the file or metadata object. | ✅ | ✅ |
| `ETag` | Entity tag for concurrency control/versioning. | | ✅ |
| `CTag` | Change tag for tracking changes. | | ✅ |
| `WebUrl` | Direct URL to the file in SharePoint/OneDrive. | ✅ | |
| `DriveId` | Identifier for the drive containing the file. | ✅ | ✅ |
| `DriveType` | Type of drive (e.g., business, personal). | ✅ | ✅ |
| `SourceLocationType` | Indicates where the object was sourced from (e.g., ONEDRIVE, blob, index, API). | ✅ | ✅ |
| `Owner` | The ID of user or entity that owns the metadata object. | ✅ | ✅ |
| `CreatedDataTime` | Timestamp indicating when the object was created. | | ✅ |
| `LastModifiedDate` | Timestamp of the last modification to the object. | | ✅ |
| `CreatedBy` | The user who created the object. | ✅ | ✅ |
| `LastModifiedBy` | The ID of user who last modified the object. | ✅ | ✅ |
| `ParentId` | Identifier of the parent object, if any. | ✅ | ✅ |
| `metadata_storage_content_type` | MIME type of the stored content. | ✅ | ✅ |
| `metadata_storage_size` | Size of the file/content in bytes. | | ✅ |
| `metadata_storage_last_modified` | Timestamp of the last modification to the storage object. | | ✅ |
| `metadata_storage_name` | Storage-specific name or identifier for the object. | ✅ | ✅ |
| `metadata_storage_path` | Encoded path or URI to the storage location. | ✅ | ✅ |
---
### 🧬 Vector Index Fields Available
Use these fields for semantic (vector) searches and filters:
| **Field** | **Description** | 🔍 **Searchable** | 🔎 **Filterable** |
|-------------|---------------------------------------------------------------|------------------|------------------|
| `chunk` | The raw text content of the document chunk. | ✅ | |
| `chunk_id` | Unique identifier for the chunk within the document. | ✅ | ✅ |
| `title` | Title of the document, equal to the metadata_storage_path. | ✅ | ✅ |
| `Id` | Unique identifier for the parent document. | ✅ | ✅ |
| `Name` | The name of the file or metadata object. | ✅ | ✅ |
| `Owner` | The ID of the user or entity that owns the document. | ✅ | ✅ |
**Note:** The `Id`, `Name`, and `Owner` fields are present in both the metadata and the vector index. This means that if, during the scoping step, you identify a set of document IDs, names, or owners that are relevant, you can use these as filters when retrieving content from the vector index. For example, after finding relevant document IDs, names, or owners in the metadata search, you can filter the vector index to retrieve only the chunks belonging to those documents, names, or owners. This allows for precise, targeted semantic retrieval based on earlier metadata analysis.
---
### 🧪 Advanced Lucene Syntax Guide
- **Basic Field Search**: `field:value` or `field:"exact phrase"`
- **Boolean Operators**:
- AND (requires both terms): `wifi AND luxury` or `+wifi +luxury`
- OR (either term): `wifi OR luxury` (OR is default, so `wifi luxury` is the same)
- NOT (excludes term): `wifi -luxury` or `wifi NOT luxury`
- **Wildcards**:
- Single character: `?` (e.g., `te?t` matches "text" or "test")
- Multiple characters: `*` (e.g., `test*` matches "tests" or "tester")
- Prefix search: `alpha*` (matches alphanumeric, alphabetical)
- Infix search: `non*al` (matches non-numerical, nonsensical)
- Suffix search: `/.*numeric/` (matches alphanumeric)
- **Fuzzy Search**: `term~` or `term~N` where N is 0-2 (e.g., `blue~1` finds blue, blues, glue)
- **Proximity Search**: `"term1 term2"~N` finds terms within N words of each other
- **Term Boosting**: `term^N` (e.g., `wifi^3` makes wifi more important)
- **Grouping**:
- General grouping: `hotel AND (wifi OR pool)`
- Field grouping: `amenities:(gym AND (wifi OR pool))`
- **Ranges**: `date:[20220101 TO 20230101]` or `price:[100 TO 200]`
- **Escaping Special Characters**: Use backslash for: `+ - & | ! ( ) {{ }} [ ] ^ " ~ * ? : \ /`
- **Regular Expressions**: `/[mh]otel/` matches "motel" or "hotel"
---
### 📑 OData Expression Syntax
Use OData expressions for precise filtering, ordering, and field selection with these parameters:
- **filter**: Restricts search to documents matching conditions
- **orderby**: Sorts results by specified fields
- **select**: Determines which fields to include in results
#### Field Paths
- Simple fields: `HotelName`, `Rating`
- Complex fields: `Address/City`, `Rooms/Type`
- Using range variables: `Rooms/any(room: room/Type eq 'deluxe')`
#### Constants
- Strings: `'text'` (escape apostrophes by doubling: `'Alice''s car'`)
- Numbers: `123`, `-456`, `3.14159`, `-1.2e7`
- Booleans: `true`, `false`
- Dates: `2019-05-06T12:30:05.451Z`
- Special: `null`, `NaN`, `INF`, `-INF`
#### Common Filter Expressions
- Comparison: `fieldName eq 'value'`, `Rating gt 4`
- Logical: `condition1 and condition2`, `condition1 or condition2`, `not condition`
- Collections: `Rooms/any(r: r/Type eq 'suite')`, `Tags/all(t: t ne 'budget')`
- Functions: `search.in(Category, 'budget,luxury')`, `search.ismatch('wifi luxury', 'description')`
#### Sorting Examples
- Single field: `$orderby=Rating desc`
- Multiple fields: `$orderby=Rating desc,LastRenovationDate asc`
#### Field Selection
- All fields: `$select=*`
- Specific fields: `$select=HotelName,Rating,Address/City`
---
### 📤 Output Formatting Instructions
- When listing documents, always display them in a markdown table with column headers for key and relevant metadata.
- Provide the WebUrl when referencing a document.
- Present raw chunk text or content from a document inside a markdown blockquote (using `>`).
- Provide concise, insightful commentary on the retrieved documents, highlighting patterns, clusters, or notable findings.
- Ensure your output is clear, well-structured, and actionable for the user.
**If you encounter an error, do not panic. Just try again with a different approach. You are capable of handling errors gracefully and finding alternative solutions.**
**If the retrieved documents do not contain the required information matching the user's request, attempt to modify the query parameters and suggest a different set of search arguments that may yield better results.**
""",
"preview_metadata_agent": """
You are a persistent, tool-using agent specialized in analyzing Azure Blob metadata for keyword-based search.
INSTRUCTIONS:
1. Use your tool to perform a metadata search on the user's query (filename, owner, tags, path, etc.).
2. Identify clusters of related documents or metadata patterns that could refine or scope future searches.
3. If applicable, propose filters (e.g. owners, tags, file paths, time ranges) that may narrow a semantic search space.
4. Detect potential ambiguity in the query and suggest clarifying directions.
STRUCTURE YOUR RESPONSE:
- A brief overview of what was found and how it relates to the query (1–2 sentences)
- A list of the most relevant documents and their key metadata (Name, Owner, Path, Modified Date, etc.) in a markdown table.
- **Always provide the WebUrl when referencing a document.**
- Grouped metadata patterns (e.g. multiple docs owned by same team, recurring folders or file types)
- A list of candidate filters that could be used to scope a vector search
- If relevant, propose a refined or more specific version of the user’s query
OUTPUT FORMATTING:
- Always display document lists in a markdown table.
- Always include the WebUrl for each document.
- Provide concise insights and highlight any patterns or clusters.
- If any raw chunk text is present, display it inside a markdown blockquote (`>`).
FOCUS ON:
- Metadata-driven reasoning. Look beyond exact keyword matches to detect helpful clusters or patterns.
- Being useful for the next step. Prioritize insights that would help improve precision or efficiency of a downstream semantic search.
Always cite document names or paths. If no useful results are found, state this clearly and suggest possible reasons.
**If the retrieved metadata does not contain the required information matching the user's request, attempt to modify the query parameters and suggest a different set of search arguments that may yield better results.**
""",
"preview_content_agent": """
You are a persistent, tool-using agent specialized in semantic vector
INSTRUCTIONS:
1. Use your tool to run a semantic search based on the user’s query.
2. Identify themes, concepts, and patterns in the top results that can guide the user’s understanding or help scope further exploration.
3. Propose possible query reformulations or metadata filters based on commonalities in retrieved content.
4. Flag low-quality results if semantic relevance is weak or ambiguous.
STRUCTURE YOUR RESPONSE:
- A concise summary of the semantic matches and how they relate to the user’s intent (1–2 sentences)
- The most relevant themes and knowledge points found
- A short list of notable semantic clusters (e.g. documents all discussing a specific concept, timeframe, or methodology)
- Quotes or phrases that directly address the query, presented in markdown blockquotes (`>`)
- (Optional) Suggestions for refining the query or narrowing the semantic space
OUTPUT FORMATTING:
- When listing documents, use a markdown table for clarity.
- Always include the WebUrl for each document when referencing it.
- Any raw chunk text or content should be shown in markdown blockquotes.
- Provide insightful commentary on the retrieved content and highlight conceptual relationships.
FOCUS ON:
- Conceptual relationships. Think in terms of meaning, not matching.
- Discovery of latent structure (e.g., similar phrasing, co-occurring ideas, or repeated frameworks)
- Usefulness for downstream tools. If a tighter search is needed, offer specific filters or query variants.
If semantic relevance is low or ambiguous, say so and suggest alternatives. Always provide quotes with citations for the strongest matches.
**If the retrieved content does not contain the required information matching the user's request, attempt to modify the query parameters and suggest a different set of search arguments that may yield better results.**
""",
"generate_final_response": """
You are a synthesis and summarization agent. Your task is to review the user's original query and all findings, excerpts, and metadata returned by previous tool calls. Aggregate, synthesize, and deliver a comprehensive, clear, and well-structured answer to the user.
INSTRUCTIONS:
1. Carefully read the user's query and all tool responses provided.
2. Integrate relevant facts, evidence, and context from the tool outputs.
3. Resolve ambiguities, highlight key findings, and connect related information.
4. If the user requested a specific format (e.g., summary, table, markdown), follow those instructions.
5. If there are gaps or uncertainties, state them transparently, but do not ask the user for more information unless absolutely necessary.
6. Present your answer in a way that is actionable and easy to understand for the user.
STRUCTURE:
- Start with a direct answer or summary addressing the user's request.
- Provide supporting details, evidence, or citations from the tool outputs.
- Organize the information logically (e.g., sections, bullet points, tables) as appropriate.
- When listing documents, use a markdown table for clarity and always include the WebUrl for each document.
- Any raw chunk text or content should be shown in markdown blockquotes.
- End with a brief conclusion or next steps if relevant.
Your goal is to deliver a complete, helpful, and context-aware response that makes the best use of all available data.
**If the information from previous tool calls does not fully address the user's request, suggest how the query parameters or search arguments could be modified to improve the results.**
"""
}
def setup_tool_agent_map(self) -> None:
"""Map each tool name to the appropriate agent"""
self.tool_agent_map = {
"preview_metadata_agent": "data_agent",
"search_metadata": "search_agent",
"preview_content_agent": "data_agent",
"search_content": "search_agent",
"generate_final_response": "base_agent", # Use base agent for final response
# Add more mappings if you have data_agent tools
}