vectara-agentic 0.1.28__py3-none-any.whl → 0.2.1__py3-none-any.whl

vectara_agentic/agent.py

@@ -1,8 +1,9 @@
 """
 This module contains the Agent class for handling different types of agents and their interactions.
 """
-from typing import List, Callable, Optional, Dict, Any
+from typing import List, Callable, Optional, Dict, Any, Union, Tuple
 import os
+import re
 from datetime import date
 import time
 import json
@@ -10,12 +11,15 @@ import logging
 import traceback
 import asyncio
 
-import dill
+import cloudpickle as pickle
+
 from dotenv import load_dotenv
 
 from retrying import retry
 from pydantic import Field, create_model
 
+from llama_index.core.memory import ChatMemoryBuffer
+from llama_index.core.llms import ChatMessage, MessageRole
 from llama_index.core.tools import FunctionTool
 from llama_index.core.agent import ReActAgent
 from llama_index.core.agent.react.formatter import ReActChatFormatter
@@ -24,7 +28,7 @@ from llama_index.agent.lats import LATSAgentWorker
 from llama_index.core.callbacks import CallbackManager, TokenCountingHandler
 from llama_index.core.callbacks.base_handler import BaseCallbackHandler
 from llama_index.agent.openai import OpenAIAgent
-from llama_index.core.memory import ChatMemoryBuffer
+
 
 from .types import AgentType, AgentStatusType, LLMRole, ToolType, AgentResponse, AgentStreamingResponse
 from .utils import get_llm, get_tokenizer_for_model
@@ -35,6 +39,21 @@ from .tools import VectaraToolFactory, VectaraTool, ToolsFactory
 from .tools_catalog import get_current_date
 from .agent_config import AgentConfig
 
+class IgnoreUnpickleableAttributeFilter(logging.Filter):
+    '''
+    Filter to ignore log messages that contain certain strings
+    '''
+    def filter(self, record):
+        msgs_to_ignore = [
+            "Removing unpickleable private attribute _chunking_tokenizer_fn",
+            "Removing unpickleable private attribute _split_fns",
+            "Removing unpickleable private attribute _sub_sentence_split_fns",
+        ]
+        return all(msg not in record.getMessage() for msg in msgs_to_ignore)
+
+
+logging.getLogger().addFilter(IgnoreUnpickleableAttributeFilter())
+
 logger = logging.getLogger("opentelemetry.exporter.otlp.proto.http.trace_exporter")
 logger.setLevel(logging.CRITICAL)
 
@@ -81,6 +100,34 @@ def _retry_if_exception(exception):
     return isinstance(exception, (TimeoutError))
 
 
+def get_field_type(field_schema: dict) -> Any:
+    """
+    Convert a JSON schema field definition to a Python type.
+    Handles 'type' and 'anyOf' cases.
+    """
+    json_type_to_python = {
+        "string": str,
+        "integer": int,
+        "boolean": bool,
+        "array": list,
+        "object": dict,
+        "number": float,
+    }
+    if "anyOf" in field_schema:
+        types = []
+        for option in field_schema["anyOf"]:
+            # If the option has a type, convert it; otherwise, use Any.
+            if "type" in option:
+                types.append(json_type_to_python.get(option["type"], Any))
+            else:
+                types.append(Any)
+        # Return a Union of the types. For example, Union[str, int]
+        return Union[tuple(types)]
+    elif "type" in field_schema:
+        return json_type_to_python.get(field_schema["type"], Any)
+    else:
+        return Any
+
 class Agent:
     """
     Agent class for handling different types of agents and their interactions.
@@ -96,6 +143,7 @@ class Agent:
         agent_progress_callback: Optional[Callable[[AgentStatusType, str], None]] = None,
         query_logging_callback: Optional[Callable[[str, str], None]] = None,
         agent_config: Optional[AgentConfig] = None,
+        chat_history: Optional[list[Tuple[str, str]]] = None,
     ) -> None:
         """
         Initialize the agent with the specified type, tools, topic, and system message.
@@ -111,10 +159,13 @@ class Agent:
             query_logging_callback (Callable): A callback function the code calls upon completion of a query
             agent_config (AgentConfig, optional): The configuration of the agent.
                 Defaults to AgentConfig(), which reads from environment variables.
+            chat_history (Tuple[str, str], optional): A list of user/agent chat pairs to initialize the agent memory.
         """
         self.agent_config = agent_config or AgentConfig()
         self.agent_type = self.agent_config.agent_type
-        self.tools = tools + [ToolsFactory().create_tool(get_current_date)]
+        self.tools = tools
+        if not any(tool.metadata.name == 'get_current_date' for tool in self.tools):
+            self.tools += [ToolsFactory().create_tool(get_current_date)]
         self.llm = get_llm(LLMRole.MAIN, config=self.agent_config)
         self._custom_instructions = custom_instructions
         self._topic = topic
@@ -135,7 +186,14 @@ class Agent:
         self.llm.callback_manager = callback_manager
         self.verbose = verbose
 
-        self.memory = ChatMemoryBuffer.from_defaults(token_limit=128000)
+        if chat_history:
+            msg_history = []
+            for inx, text in enumerate(chat_history):
+                role = MessageRole.USER if inx % 2 == 0 else MessageRole.ASSISTANT
+                msg_history.append(ChatMessage.from_str(content=text, role=role))
+            self.memory = ChatMemoryBuffer.from_defaults(token_limit=128000, chat_history=msg_history)
+        else:
+            self.memory = ChatMemoryBuffer.from_defaults(token_limit=128000)
         if self.agent_type == AgentType.REACT:
             prompt = _get_prompt(REACT_PROMPT_TEMPLATE, topic, custom_instructions)
             self.agent = ReActAgent.from_tools(
@@ -219,7 +277,10 @@ class Agent:
 
         # Compare tools
         if self.tools != other.tools:
-            print(f"Comparison failed: tools differ. (self.tools: {self.tools}, other.tools: {other.tools})")
+            print(
+                "Comparison failed: tools differ."
+                f"(self.tools: {[t.metadata.name for t in self.tools]}, "
+                f"other.tools: {[t.metadata.name for t in other.tools]})")
             return False
 
         # Compare topic
@@ -263,6 +324,7 @@ class Agent:
         agent_progress_callback: Optional[Callable[[AgentStatusType, str], None]] = None,
         query_logging_callback: Optional[Callable[[str, str], None]] = None,
         agent_config: AgentConfig = AgentConfig(),
+        chat_history: Optional[list[Tuple[str, str]]] = None,
     ) -> "Agent":
         """
         Create an agent from tools, agent type, and language model.
@@ -277,6 +339,7 @@ class Agent:
                 update_func (Callable): old name for agent_progress_callback. Will be deprecated in future.
             query_logging_callback (Callable): A callback function the code calls upon completion of a query
             agent_config (AgentConfig, optional): The configuration of the agent.
+            chat_history (Tuple[str, str], optional): A list of user/agent chat pairs to initialize the agent memory.
 
         Returns:
             Agent: An instance of the Agent class.
@@ -285,7 +348,8 @@ class Agent:
             tools=tools, topic=topic, custom_instructions=custom_instructions,
             verbose=verbose, agent_progress_callback=agent_progress_callback,
             query_logging_callback=query_logging_callback,
-            update_func=update_func, agent_config=agent_config
+            update_func=update_func, agent_config=agent_config,
+            chat_history=chat_history,
         )
 
     @classmethod
@@ -294,28 +358,42 @@ class Agent:
         tool_name: str,
         data_description: str,
         assistant_specialty: str,
-        vectara_customer_id: str = str(os.environ.get("VECTARA_CUSTOMER_ID", "")),
-        vectara_corpus_id: str = str(os.environ.get("VECTARA_CORPUS_ID", "")),
+        vectara_corpus_key: str = str(os.environ.get("VECTARA_CORPUS_KEY", "")),
         vectara_api_key: str = str(os.environ.get("VECTARA_API_KEY", "")),
         agent_progress_callback: Optional[Callable[[AgentStatusType, str], None]] = None,
         query_logging_callback: Optional[Callable[[str, str], None]] = None,
         verbose: bool = False,
         vectara_filter_fields: list[dict] = [],
+        vectara_offset: int = 0,
         vectara_lambda_val: float = 0.005,
-        vectara_reranker: str = "mmr",
+        vectara_semantics: str = "default",
+        vectara_custom_dimensions: Dict = {},
+        vectara_reranker: str = "slingshot",
         vectara_rerank_k: int = 50,
+        vectara_rerank_limit: Optional[int] = None,
+        vectara_rerank_cutoff: Optional[float] = None,
+        vectara_diversity_bias: float = 0.2,
+        vectara_udf_expression: str = None,
+        vectara_rerank_chain: List[Dict] = None,
         vectara_n_sentences_before: int = 2,
         vectara_n_sentences_after: int = 2,
         vectara_summary_num_results: int = 10,
-        vectara_summarizer: str = "vectara-summary-ext-24-05-sml",
+        vectara_summarizer: str = "vectara-summary-ext-24-05-med-omni",
+        vectara_summary_response_language: str = "eng",
+        vectara_summary_prompt_text: Optional[str] = None,
+        vectara_max_response_chars: Optional[int] = None,
+        vectara_max_tokens: Optional[int] = None,
+        vectara_temperature: Optional[float] = None,
+        vectara_frequency_penalty: Optional[float] = None,
+        vectara_presence_penalty: Optional[float] = None,
+        vectara_save_history: bool = True,
     ) -> "Agent":
         """
         Create an agent from a single Vectara corpus
 
         Args:
             tool_name (str): The name of Vectara tool used by the agent
-            vectara_customer_id (str): The Vectara customer ID.
-            vectara_corpus_id (str): The Vectara corpus ID (or comma separated list of IDs).
+            vectara_corpus_key (str): The Vectara corpus key (or comma separated list of keys).
             vectara_api_key (str): The Vectara API key.
             agent_progress_callback (Callable): A callback function the code calls on any agent updates.
             query_logging_callback (Callable): A callback function the code calls upon completion of a query
@@ -324,21 +402,41 @@ class Agent:
             verbose (bool, optional): Whether to print verbose output.
             vectara_filter_fields (List[dict], optional): The filterable attributes
                 (each dict maps field name to Tuple[type, description]).
-            vectara_lambda_val (float, optional): The lambda value for Vectara hybrid search.
-            vectara_reranker (str, optional): The Vectara reranker name (default "mmr")
+            vectara_offset (int, optional): Number of results to skip.
+            vectara_lambda_val (float, optional): Lambda value for Vectara hybrid search.
+            vectara_semantics: (str, optional): Indicates whether the query is intended as a query or response.
+            vectara_custom_dimensions: (Dict, optional): Custom dimensions for the query.
+            vectara_reranker (str, optional): The Vectara reranker name (default "slingshot")
             vectara_rerank_k (int, optional): The number of results to use with reranking.
+            vetara_rerank_limit: (int, optional): The maximum number of results to return after reranking.
+            vectara_rerank_cutoff: (float, optional): The minimum score threshold for results to include after
+                reranking.
+            vectara_diversity_bias (float, optional): The MMR diversity bias.
+            vectara_udf_expression (str, optional): The user defined expression for reranking results.
+            vectara_rerank_chain (List[Dict], optional): A list of Vectara rerankers to be applied sequentially.
             vectara_n_sentences_before (int, optional): The number of sentences before the matching text
             vectara_n_sentences_after (int, optional): The number of sentences after the matching text.
             vectara_summary_num_results (int, optional): The number of results to use in summarization.
             vectara_summarizer (str, optional): The Vectara summarizer name.
+            vectara_summary_response_language (str, optional): The response language for the Vectara summary.
+            vectara_summary_prompt_text (str, optional): The custom prompt, using appropriate prompt variables and
+                functions.
+            vectara_max_response_chars (int, optional): The desired maximum number of characters for the generated
+                summary.
+            vectara_max_tokens (int, optional): The maximum number of tokens to be returned by the LLM.
+            vectara_temperature (float, optional): The sampling temperature; higher values lead to more randomness.
+            vectara_frequency_penalty (float, optional): How much to penalize repeating tokens in the response,
+                higher values reducing likelihood of repeating the same line.
+            vectara_presence_penalty (float, optional): How much to penalize repeating tokens in the response,
+                higher values increasing the diversity of topics.
+            vectara_save_history (bool, optional): Whether to save the query in history.
 
         Returns:
             Agent: An instance of the Agent class.
         """
         vec_factory = VectaraToolFactory(
             vectara_api_key=vectara_api_key,
-            vectara_customer_id=vectara_customer_id,
-            vectara_corpus_id=vectara_corpus_id,
+            vectara_corpus_key=vectara_corpus_key,
         )
         field_definitions = {}
         field_definitions["query"] = (str, Field(description="The user query"))  # type: ignore
@@ -349,8 +447,12 @@ class Agent:
             )  # type: ignore
         query_args = create_model("QueryArgs", **field_definitions)  # type: ignore
 
+        # tool name must be valid Python function name
+        if tool_name:
+            tool_name = re.sub(r"[^A-Za-z0-9_]", "_", tool_name)
+
         vectara_tool = vec_factory.create_rag_tool(
-            tool_name=tool_name or f"vectara_{vectara_corpus_id}",
+            tool_name=tool_name or f"vectara_{vectara_corpus_key}",
             tool_description=f"""
             Given a user query,
             returns a response (str) to a user question about {data_description}.
@@ -358,11 +460,27 @@ class Agent:
             tool_args_schema=query_args,
             reranker=vectara_reranker,
             rerank_k=vectara_rerank_k,
+            rerank_limit=vectara_rerank_limit,
+            rerank_cutoff=vectara_rerank_cutoff,
+            mmr_diversity_bias=vectara_diversity_bias,
+            udf_expression=vectara_udf_expression,
+            rerank_chain=vectara_rerank_chain,
             n_sentences_before=vectara_n_sentences_before,
             n_sentences_after=vectara_n_sentences_after,
+            offset=vectara_offset,
             lambda_val=vectara_lambda_val,
+            semantics=vectara_semantics,
+            custom_dimensions=vectara_custom_dimensions,
             summary_num_results=vectara_summary_num_results,
             vectara_summarizer=vectara_summarizer,
+            summary_response_lang=vectara_summary_response_language,
+            vectara_prompt_text=vectara_summary_prompt_text,
+            max_response_chars=vectara_max_response_chars,
+            max_tokens=vectara_max_tokens,
+            temperature=vectara_temperature,
+            frequency_penalty=vectara_frequency_penalty,
+            presence_penalty=vectara_presence_penalty,
+            save_history=vectara_save_history,
             include_citations=True,
             verbose=verbose,
         )
@@ -534,12 +652,13 @@ class Agent:
 
         for tool in self.tools:
             # Serialize each tool's metadata, function, and dynamic model schema (QueryArgs)
+            # TODO: deal with tools that have weakref (e.g. db_tools); for now those cannot be serialized.
             tool_dict = {
                 "tool_type": tool.metadata.tool_type.value,
                 "name": tool.metadata.name,
                 "description": tool.metadata.description,
-                "fn": dill.dumps(tool.fn).decode("latin-1") if tool.fn else None,  # Serialize fn
-                "async_fn": dill.dumps(tool.async_fn).decode("latin-1")
+                "fn": pickle.dumps(tool.fn).decode("latin-1") if tool.fn else None,  # Serialize fn
+                "async_fn": pickle.dumps(tool.async_fn).decode("latin-1")
                 if tool.async_fn
                 else None,  # Serialize async_fn
                 "fn_schema": tool.metadata.fn_schema.model_json_schema()
@@ -550,7 +669,7 @@ class Agent:
 
         return {
             "agent_type": self.agent_type.value,
-            "memory": dill.dumps(self.agent.memory).decode("latin-1"),
+            "memory": pickle.dumps(self.agent.memory).decode("latin-1"),
             "tools": tool_info,
             "topic": self._topic,
             "custom_instructions": self._custom_instructions,
@@ -564,39 +683,30 @@ class Agent:
         agent_config = AgentConfig.from_dict(data["agent_config"])
         tools = []
 
-        json_type_to_python = {
-            "string": str,
-            "integer": int,
-            "boolean": bool,
-            "array": list,
-            "object": dict,
-            "number": float,
-        }
-
         for tool_data in data["tools"]:
             # Recreate the dynamic model using the schema info
             if tool_data.get("fn_schema"):
                 field_definitions = {}
                 for field, values in tool_data["fn_schema"]["properties"].items():
+                    # Instead of checking for 'type', use the helper:
+                    field_type = get_field_type(values)
+                    # If there's a default value, include it.
                     if "default" in values:
                         field_definitions[field] = (
-                            json_type_to_python.get(values["type"], values["type"]),
-                            Field(
-                                description=values["description"],
-                                default=values["default"],
-                            ),
-                        )  # type: ignore
+                            field_type,
+                            Field(description=values.get("description", ""), default=values["default"]),
+                        )
                     else:
                         field_definitions[field] = (
-                            json_type_to_python.get(values["type"], values["type"]),
-                            Field(description=values["description"]),
-                        )  # type: ignore
+                            field_type,
+                            Field(description=values.get("description", "")),
+                        )
                 query_args_model = create_model("QueryArgs", **field_definitions)  # type: ignore
             else:
                 query_args_model = create_model("QueryArgs")
 
-            fn = dill.loads(tool_data["fn"].encode("latin-1")) if tool_data["fn"] else None
-            async_fn = dill.loads(tool_data["async_fn"].encode("latin-1")) if tool_data["async_fn"] else None
+            fn = pickle.loads(tool_data["fn"].encode("latin-1")) if tool_data["fn"] else None
+            async_fn = pickle.loads(tool_data["async_fn"].encode("latin-1")) if tool_data["async_fn"] else None
 
             tool = VectaraTool.from_defaults(
                 name=tool_data["name"],
@@ -615,7 +725,7 @@ class Agent:
             custom_instructions=data["custom_instructions"],
             verbose=data["verbose"],
         )
-        memory = dill.loads(data["memory"].encode("latin-1")) if data.get("memory") else None
+        memory = pickle.loads(data["memory"].encode("latin-1")) if data.get("memory") else None
         if memory:
             agent.agent.memory = memory
         return agent

vectara_agentic/agent_config.py

@@ -44,6 +44,15 @@ class AgentConfig:
         default_factory=lambda: os.getenv("VECTARA_AGENTIC_TOOL_MODEL_NAME", "")
     )
 
+    # Params for Private LLM endpoint if used
+    private_llm_api_base: str = field(
+        default_factory=lambda: os.getenv("VECTARA_AGENTIC_PRIVATE_LLM_API_BASE",
+                                          "http://private-endpoint.company.com:5000/v1")
+    )
+    private_llm_api_key: str = field(
+        default_factory=lambda: os.getenv("VECTARA_AGENTIC_PRIVATE_LLM_API_KEY", "<private-api-key>")
+    )
+
     # Observer
     observer: ObserverType = field(
         default_factory=lambda: ObserverType(

vectara_agentic/db_tools.py

@@ -74,7 +74,7 @@ class DBLoadUniqueValues(DBTool):
     """
     A tool to list all unique values for each column in a set of columns of a database table.
     """
-    def __call__(self, table_name: str, columns: list[str], num_vals: int = 200) -> dict:
+    def __call__(self, table_name: str, columns: list[str], num_vals: int = 200) -> Any:
         """
         Fetches the first num_vals unique values from the specified columns of the database table.
 
@@ -84,7 +84,7 @@ class DBLoadUniqueValues(DBTool):
             num_vals (int): The number of unique values to fetch for each column. Default is 200.
 
         Returns:
-            dict: A dictionary containing the unique values for each column.
+            Any: the result of the database query
         """
         res = {}
         try: