vectara-agentic 0.2.7__py3-none-any.whl → 0.2.8__py3-none-any.whl

tests/test_agent.py

@@ -106,43 +106,6 @@ class TestAgentPackage(unittest.TestCase):
         self.assertIsInstance(agent, Agent)
         self.assertEqual(agent._topic, "question answering")
 
-    def test_serialization(self):
-        with ARIZE_LOCK:
-            config = AgentConfig(
-                agent_type=AgentType.REACT,
-                main_llm_provider=ModelProvider.ANTHROPIC,
-                tool_llm_provider=ModelProvider.TOGETHER,
-                observer=ObserverType.ARIZE_PHOENIX
-            )
-
-            agent = Agent.from_corpus(
-                tool_name="RAG Tool",
-                agent_config=config,
-                vectara_corpus_key="corpus_key",
-                vectara_api_key="api_key",
-                data_description="information",
-                assistant_specialty="question answering",
-            )
-
-            agent_reloaded = agent.loads(agent.dumps())
-            agent_reloaded_again = agent_reloaded.loads(agent_reloaded.dumps())
-
-            self.assertIsInstance(agent_reloaded, Agent)
-            self.assertEqual(agent, agent_reloaded)
-            self.assertEqual(agent.agent_type, agent_reloaded.agent_type)
-
-            self.assertEqual(agent.agent_config.observer, agent_reloaded.agent_config.observer)
-            self.assertEqual(agent.agent_config.main_llm_provider, agent_reloaded.agent_config.main_llm_provider)
-            self.assertEqual(agent.agent_config.tool_llm_provider, agent_reloaded.agent_config.tool_llm_provider)
-
-            self.assertIsInstance(agent_reloaded, Agent)
-            self.assertEqual(agent, agent_reloaded_again)
-            self.assertEqual(agent.agent_type, agent_reloaded_again.agent_type)
-
-            self.assertEqual(agent.agent_config.observer, agent_reloaded_again.agent_config.observer)
-            self.assertEqual(agent.agent_config.main_llm_provider, agent_reloaded_again.agent_config.main_llm_provider)
-            self.assertEqual(agent.agent_config.tool_llm_provider, agent_reloaded_again.agent_config.tool_llm_provider)
-
     def test_chat_history(self):
         tools = [ToolsFactory().create_tool(mult)]
         topic = "AI topic"

tests/test_serialization.py

@@ -0,0 +1,110 @@
+import unittest
+import threading
+import os
+
+from vectara_agentic.agent import Agent, AgentType
+from vectara_agentic.agent_config import AgentConfig
+from vectara_agentic.types import ModelProvider, ObserverType
+from vectara_agentic.tools import ToolsFactory
+
+from llama_index.core.utilities.sql_wrapper import SQLDatabase
+from sqlalchemy import create_engine
+
+def mult(x: float, y: float) -> float:
+    return x * y
+
+
+ARIZE_LOCK = threading.Lock()
+
+class TestAgentSerialization(unittest.TestCase):
+
+    @classmethod
+    def tearDown(cls):
+        try:
+            os.remove('ev_database.db')
+        except FileNotFoundError:
+            pass
+
+    def test_serialization(self):
+        with ARIZE_LOCK:
+            config = AgentConfig(
+                agent_type=AgentType.REACT,
+                main_llm_provider=ModelProvider.ANTHROPIC,
+                tool_llm_provider=ModelProvider.TOGETHER,
+                observer=ObserverType.ARIZE_PHOENIX
+            )
+            db_tools = ToolsFactory().database_tools(
+                tool_name_prefix = "ev",
+                content_description = 'Electric Vehicles in the state of Washington and other population information',
+                sql_database = SQLDatabase(create_engine('sqlite:///ev_database.db')),
+            )
+
+            tools = [ToolsFactory().create_tool(mult)] + ToolsFactory().standard_tools() + db_tools
+            topic = "AI topic"
+            instructions = "Always do as your father tells you, if your mother agrees!"
+            agent = Agent(
+                tools=tools,
+                topic=topic,
+                custom_instructions=instructions,
+                agent_config=config
+            )
+
+            agent_reloaded = agent.loads(agent.dumps())
+            agent_reloaded_again = agent_reloaded.loads(agent_reloaded.dumps())
+
+            self.assertIsInstance(agent_reloaded, Agent)
+            self.assertEqual(agent, agent_reloaded)
+            self.assertEqual(agent.agent_type, agent_reloaded.agent_type)
+
+            self.assertEqual(agent.agent_config.observer, agent_reloaded.agent_config.observer)
+            self.assertEqual(agent.agent_config.main_llm_provider, agent_reloaded.agent_config.main_llm_provider)
+            self.assertEqual(agent.agent_config.tool_llm_provider, agent_reloaded.agent_config.tool_llm_provider)
+
+            self.assertIsInstance(agent_reloaded, Agent)
+            self.assertEqual(agent, agent_reloaded_again)
+            self.assertEqual(agent.agent_type, agent_reloaded_again.agent_type)
+
+            self.assertEqual(agent.agent_config.observer, agent_reloaded_again.agent_config.observer)
+            self.assertEqual(agent.agent_config.main_llm_provider, agent_reloaded_again.agent_config.main_llm_provider)
+            self.assertEqual(agent.agent_config.tool_llm_provider, agent_reloaded_again.agent_config.tool_llm_provider)
+
+    def test_serialization_from_corpus(self):
+        with ARIZE_LOCK:
+            config = AgentConfig(
+                agent_type=AgentType.REACT,
+                main_llm_provider=ModelProvider.ANTHROPIC,
+                tool_llm_provider=ModelProvider.TOGETHER,
+                observer=ObserverType.ARIZE_PHOENIX
+            )
+
+            agent = Agent.from_corpus(
+                tool_name="RAG Tool",
+                agent_config=config,
+                vectara_corpus_key="corpus_key",
+                vectara_api_key="api_key",
+                data_description="information",
+                assistant_specialty="question answering",
+            )
+
+            agent_reloaded = agent.loads(agent.dumps())
+            agent_reloaded_again = agent_reloaded.loads(agent_reloaded.dumps())
+
+            self.assertIsInstance(agent_reloaded, Agent)
+            self.assertEqual(agent, agent_reloaded)
+            self.assertEqual(agent.agent_type, agent_reloaded.agent_type)
+
+            self.assertEqual(agent.agent_config.observer, agent_reloaded.agent_config.observer)
+            self.assertEqual(agent.agent_config.main_llm_provider, agent_reloaded.agent_config.main_llm_provider)
+            self.assertEqual(agent.agent_config.tool_llm_provider, agent_reloaded.agent_config.tool_llm_provider)
+
+            self.assertIsInstance(agent_reloaded, Agent)
+            self.assertEqual(agent, agent_reloaded_again)
+            self.assertEqual(agent.agent_type, agent_reloaded_again.agent_type)
+
+            self.assertEqual(agent.agent_config.observer, agent_reloaded_again.agent_config.observer)
+            self.assertEqual(agent.agent_config.main_llm_provider, agent_reloaded_again.agent_config.main_llm_provider)
+            self.assertEqual(agent.agent_config.tool_llm_provider, agent_reloaded_again.agent_config.tool_llm_provider)
+
+
+if __name__ == "__main__":
+    unittest.main()

vectara_agentic/_version.py

@@ -1,4 +1,4 @@
 """
 Define the version of the package.
 """
-__version__ = "0.2.7"
+__version__ = "0.2.8"

vectara_agentic/agent.py

@@ -768,6 +768,7 @@ class Agent:
         """
         max_attempts = 4 if self.fallback_agent_config else 2
         attempt = 0
+        orig_llm = self.llm.metadata.model_name
         while attempt < max_attempts:
             try:
                 current_agent = self._get_current_agent()
@@ -788,16 +789,20 @@ class Agent:
                 agent_response.async_response_gen = _stream_response_wrapper  # Override the generator
                 return agent_response
 
-            except Exception:
+            except Exception as e:
+                last_error = e
                 if attempt >= 2:
                     if self.verbose:
-                        print("LLM call failed. Switching agent configuration.")
+                        print(f"LLM call failed on attempt {attempt}. Switching agent configuration.")
                     self._switch_agent_config()
                 time.sleep(1)
                 attempt += 1
 
         return AgentResponse(
-            response=f"LLM failure can't be resolved after {max_attempts} attempts."
+            response=(
+                f"For {orig_llm} LLM - failure can't be resolved after "
+                f"{max_attempts} attempts ({last_error})."
+            )
         )
 
     #
@@ -861,7 +866,6 @@ class Agent:
     def to_dict(self) -> Dict[str, Any]:
         """Serialize the Agent instance to a dictionary."""
         tool_info = []
-
         for tool in self.tools:
             if hasattr(tool.metadata, "fn_schema"):
                 fn_schema_cls = tool.metadata.fn_schema

vectara_agentic/db_tools.py

@@ -1,38 +1,122 @@
 """
-This module contains the code to extend and improve DatabaseToolSpec
-Specifically adding load_sample_data and load_unique_values methods, as well as
-making sure the load_data method returns a list of text values from the database, not Document[] objects.
+This module contains the code adapted from DatabaseToolSpec
+It makes the following adjustments:
+* Adds load_sample_data and load_unique_values methods.
+* Fixes serialization.
+* Makes sure the load_data method returns a list of text values from the database (and not Document[] objects).
+* Limits the returned rows to self.max_rows.
 """
-from abc import ABC
-from typing import Callable, Any
+from typing import Any, Optional, List, Awaitable, Callable
+import asyncio
+from inspect import signature
 
-#
-# Additional database tool
-#
-class DBTool(ABC):
-    """
-    A base class for vectara-agentic database tools extensions
+from sqlalchemy import MetaData, text
+from sqlalchemy.engine import Engine
+from sqlalchemy.exc import NoSuchTableError
+from sqlalchemy.schema import CreateTable
+
+from llama_index.core.readers.base import BaseReader
+from llama_index.core.utilities.sql_wrapper import SQLDatabase
+from llama_index.core.schema import Document
+from llama_index.core.tools.function_tool import FunctionTool
+from llama_index.core.tools.types import ToolMetadata
+from llama_index.core.tools.utils import create_schema_from_function
+
+
+AsyncCallable = Callable[..., Awaitable[Any]]
+
+class DatabaseTools(BaseReader):
+    """Database tools for vectara-agentic
+    This class provides a set of tools to interact with a database.
+    It allows you to load data, list tables, describe tables, and load unique values.
+    It also provides a method to load sample data from a specified table.
     """
-    def __init__(self, load_data_fn: Callable, max_rows: int = 1000):
-        self.load_data_fn = load_data_fn
+    spec_functions = [
+        "load_data", "load_sample_data", "list_tables",
+        "describe_tables", "load_unique_values",
+    ]
+
+    def __init__(
+        self,
+        *args: Any,
+        max_rows: int = 1000,
+        sql_database: Optional[SQLDatabase] = None,
+        engine: Optional[Engine] = None,
+        uri: Optional[str] = None,
+        scheme: Optional[str] = None,
+        host: Optional[str] = None,
+        port: Optional[str] = None,
+        user: Optional[str] = None,
+        password: Optional[str] = None,
+        dbname: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
         self.max_rows = max_rows
 
-class DBLoadData(DBTool):
-    """
-    A tool to Run SQL query on the database and return the result.
-    """
-    def __call__(self, query: str) -> Any:
-        """Query and load data from the Database, returning a list of Documents.
+        if sql_database:
+            self.sql_database = sql_database
+        elif engine:
+            self.sql_database = SQLDatabase(engine, *args, **kwargs)
+        elif uri:
+            self.uri = uri
+            self.sql_database = SQLDatabase.from_uri(uri, *args, **kwargs)
+        elif (scheme and host and port and user and password and dbname):
+            uri = f"{scheme}://{user}:{password}@{host}:{port}/{dbname}"
+            self.uri = uri
+            self.sql_database = SQLDatabase.from_uri(uri, *args, **kwargs)
+        else:
+            raise ValueError(
+                "You must provide either a SQLDatabase, "
+                "a SQL Alchemy Engine, a valid connection URI, or a valid "
+                "set of credentials."
+            )
+        self._uri = getattr(self, "uri", None) or str(self.sql_database.engine.url)
+        self._metadata = MetaData()
+        self._metadata.reflect(bind=self.sql_database.engine)
+
+    def _get_metadata_from_fn_name(
+        self, fn_name: Callable,
+    ) -> Optional[ToolMetadata]:
+        """Return map from function name.
+
+        Return type is Optional, meaning that the schema can be None.
+        In this case, it's up to the downstream tool implementation to infer the schema.
+        """
+        try:
+            func = getattr(self, fn_name)
+        except AttributeError:
+            return None
+        name = fn_name
+        docstring = func.__doc__ or ""
+        description = f"{name}{signature(func)}\n{docstring}"
+        fn_schema = create_schema_from_function(fn_name, getattr(self, fn_name))
+        return ToolMetadata(name=name, description=description, fn_schema=fn_schema)
+
+    def _load_data(self, query: str) -> List[Document]:
+        documents = []
+        with self.sql_database.engine.connect() as connection:
+            if query is None:
+                raise ValueError("A query parameter is necessary to filter the data")
+            result = connection.execute(text(query))
+            for item in result.fetchall():
+                doc_str = ", ".join([str(entry) for entry in item])
+                documents.append(Document(text=doc_str))
+        return documents
 
+    def load_data(self, *args: Any, **load_kwargs: Any) -> List[str]:
+        """Query and load data from the Database, returning a list of Documents.
         Args:
             query (str): an SQL query to filter tables and rows.
-
         Returns:
-            List[text]: a list of text values from the database.
+            List[Document]: a list of Document objects from the database.
         """
+        query = args[0] if args else load_kwargs.get("args",{}).get("query")
+        if query is None:
+            raise ValueError("A query parameter is necessary to filter the data")
+
         count_query = f"SELECT COUNT(*) FROM ({query})"
         try:
-            count_rows = self.load_data_fn(count_query)
+            count_rows = self._load_data(count_query)
         except Exception as e:
             return [f"Error ({str(e)}) occurred while counting number of rows"]
         num_rows = int(count_rows[0].text)
@@ -42,19 +126,12 @@ class DBLoadData(DBTool):
                 "Please refactor your query to make it return less rows. "
             ]
         try:
-            res = self.load_data_fn(query)
+            res = self._load_data(query)
         except Exception as e:
             return [f"Error ({str(e)}) occurred while executing the query {query}"]
         return [d.text for d in res]
 
-class DBLoadSampleData(DBTool):
-    """
-    A tool to load a sample of data from the specified database table.
-
-    This tool fetches the first num_rows (default 25) rows from the given table
-    using a provided database query function.
-    """
-    def __call__(self, table_name: str, num_rows: int = 25) -> Any:
+    def load_sample_data(self, table_name: str, num_rows: int = 25) -> Any:
         """
         Fetches the first num_rows rows from the specified database table.
 
@@ -65,16 +142,39 @@ class DBLoadSampleData(DBTool):
             Any: The result of the database query.
         """
         try:
-            res = self.load_data_fn(f"SELECT * FROM {table_name} LIMIT {num_rows}")
+            res = self._load_data(f"SELECT * FROM {table_name} LIMIT {num_rows}")
         except Exception as e:
             return [f"Error ({str(e)}) occurred while loading sample data for table {table_name}"]
-        return res
+        return [d.text for d in res]
 
-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) -> Any:
+    def list_tables(self) -> List[str]:
+        """List all tables in the database.
+        Returns:
+            List[str]: A list of table names in the database.
+        """
+        return [x.name for x in self._metadata.sorted_tables]
+
+    def describe_tables(self, tables: Optional[List[str]] = None) -> str:
+        """Describe the tables in the database.
+        Args:
+            tables (Optional[List[str]]): A list of table names to describe. If None, all tables are described.
+        Returns:
+            str: A string representation of the table schemas.
+        """
+        table_names = tables or [table.name for table in self._metadata.sorted_tables]
+        table_schemas = []
+        for table_name in table_names:
+            table = next(
+                (table for table in self._metadata.sorted_tables if table.name == table_name),
+                None,
+            )
+            if table is None:
+                raise NoSuchTableError(f"Table '{table_name}' does not exist.")
+            schema = str(CreateTable(table).compile(self.sql_database.engine))
+            table_schemas.append(f"{schema}\n")
+        return "\n".join(table_schemas)
+
+    def load_unique_values(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.
 
@@ -89,8 +189,74 @@ class DBLoadUniqueValues(DBTool):
         res = {}
         try:
             for column in columns:
-                unique_vals = self.load_data_fn(f'SELECT DISTINCT "{column}" FROM {table_name} LIMIT {num_vals}')
+                unique_vals = self._load_data(f'SELECT DISTINCT "{column}" FROM {table_name} LIMIT {num_vals}')
                 res[column] = [d.text for d in unique_vals]
         except Exception as e:
             return {f"Error ({str(e)}) occurred while loading unique values for table {table_name}"}
         return res
+
+    def to_tool_list(self) -> List[FunctionTool]:
+        """
+        Returns a list of tools available.
+        """
+
+        tool_list = []
+        for tool_name in self.spec_functions:
+            func_sync = None
+            func_async = None
+            func = getattr(self, tool_name)
+            if asyncio.iscoroutinefunction(func):
+                func_async = func
+            else:
+                func_sync = func
+            metadata = self._get_metadata_from_fn_name(tool_name)
+
+            if func_sync is None:
+                if func_async is not None:
+                    func_sync = patch_sync(func_async)
+                else:
+                    raise ValueError(
+                        f"Could not retrieve a function for spec: {tool_name}"
+                    )
+
+            tool = FunctionTool.from_defaults(
+                fn=func_sync,
+                async_fn=func_async,
+                tool_metadata=metadata,
+            )
+            tool_list.append(tool)
+        return tool_list
+
+    # Custom pickling: exclude unpickleable objects
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        if "sql_database" in state:
+            state["sql_database_state"] = {"uri": self._uri}
+            del state["sql_database"]
+        if "_metadata" in state:
+            del state["_metadata"]
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        # Reconstruct the sql_database if it was removed
+        if "sql_database_state" in state:
+            uri = state["sql_database_state"].get("uri")
+            if uri:
+                self.sql_database = SQLDatabase.from_uri(uri)
+                self._uri = uri
+            else:
+                raise ValueError("Cannot reconstruct SQLDatabase without URI")
+            # Rebuild metadata after restoring the engine
+            self._metadata = MetaData()
+            self._metadata.reflect(bind=self.sql_database.engine)
+
+
+def patch_sync(func_async: AsyncCallable) -> Callable:
+    """Patch sync function from async function."""
+
+    def patched_sync(*args: Any, **kwargs: Any) -> Any:
+        loop = asyncio.get_event_loop()
+        return loop.run_until_complete(func_async(*args, **kwargs))
+
+    return patched_sync

vectara_agentic/sub_query_workflow.py

@@ -102,7 +102,7 @@ class SubQuestionQueryWorkflow(Workflow):
             - What is the name of the mayor of San Jose?
             Here is the user question: {await ctx.get('original_query')}.
             Here are previous chat messages: {chat_history}.
-            And here is the list of tools: {await ctx.get('tools')}
+            And here is the list of tools: {ev.tools}
             """,
         )
 
@@ -236,6 +236,7 @@ class SequentialSubQuestionsWorkflow(Workflow):
             print(f"Query is {await ctx.get('original_query')}")
 
         llm = await ctx.get("llm")
+        orig_query = await ctx.get("original_query")
         response = llm.complete(
             f"""
             Given a user question, and a list of tools, output a list of
@@ -256,9 +257,9 @@ class SequentialSubQuestionsWorkflow(Workflow):
             - Who is the mayor of this city?
             The answer to the first question is San Jose, which is given as context to the second question.
             The answer to the second question is Matt Mahan.
-            Here is the user question: {await ctx.get('original_query')}.
+            Here is the user question: {orig_query}.
             Here are previous chat messages: {chat_history}.
-            And here is the list of tools: {await ctx.get('tools')}
+            And here is the list of tools: {ev.tools}
             """,
         )
 

vectara_agentic/tools.py

@@ -6,6 +6,7 @@ import inspect
 import re
 import importlib
 import os
+import asyncio
 
 from typing import Callable, List, Dict, Any, Optional, Union, Type
 from pydantic import BaseModel, Field, create_model
@@ -20,8 +21,8 @@ from llama_index.core.workflow.context import Context
 
 from .types import ToolType
 from .tools_catalog import ToolsCatalog, get_bad_topics
-from .db_tools import DBLoadSampleData, DBLoadUniqueValues, DBLoadData
-from .utils import is_float, summarize_vectara_document
+from .db_tools import DatabaseTools
+from .utils import is_float, summarize_documents
 from .agent_config import AgentConfig
 
 LI_packages = {
@@ -31,7 +32,6 @@ LI_packages = {
     "exa": ToolType.QUERY,
     "neo4j": ToolType.QUERY,
     "kuzu": ToolType.QUERY,
-    "database": ToolType.QUERY,
     "google": {
         "GmailToolSpec": {
             "load_data": ToolType.QUERY,
@@ -109,9 +109,20 @@ class VectaraTool(FunctionTool):
             fn, name, description, return_direct, fn_schema, async_fn, tool_metadata,
             callback, async_callback
         )
-        vectara_tool = cls(tool_type=tool_type, fn=tool.fn, metadata=tool.metadata, async_fn=tool.async_fn)
+        vectara_tool = cls(
+            tool_type=tool_type, fn=tool.fn, metadata=tool.metadata, async_fn=tool.async_fn,
+        )
         return vectara_tool
 
+    def __str__(self) -> str:
+        return (
+            f"Tool(name={self.metadata.name}, "
+            f"Tool metadata={self.metadata})"
+        )
+
+    def __repr__(self) -> str:
+        return str(self)
+
     def __eq__(self, other):
         if not isinstance(other, VectaraTool):
             return False
@@ -508,17 +519,29 @@ class VectaraToolFactory:
                     raw_input={"args": args, "kwargs": kwargs},
                     raw_output={"response": msg},
                 )
-            tool_output = "Matching documents:\n"
             unique_ids = set()
+            docs = []
             for doc in response:
                 if doc.id_ in unique_ids:
                     continue
                 unique_ids.add(doc.id_)
-                if summarize:
-                    summary = summarize_vectara_document(self.vectara_corpus_key, self.vectara_api_key, doc.id_)
-                    tool_output += f"document_id: '{doc.id_}'\nmetadata: '{doc.metadata}'\nsummary: '{summary}'\n\n"
-                else:
+                docs.append((doc.id_, doc.metadata))
+            tool_output = "Matching documents:\n"
+            if summarize:
+                summaries_dict = asyncio.run(
+                    summarize_documents(
+                        self.vectara_corpus_key,
+                        self.vectara_api_key,
+                        list(unique_ids)
+                    )
+                )
+                for doc_id, metadata in docs:
+                    summary = summaries_dict.get(doc_id, "")
+                    tool_output += f"document_id: '{doc_id}'\nmetadata: '{metadata}'\nsummary: '{summary}'\n\n"
+            else:
+                for doc in docs:
                     tool_output += f"document_id: '{doc.id_}'\nmetadata: '{doc.metadata}'\n\n"
+
             out = ToolOutput(
                 tool_name=search_function.__name__,
                 content=tool_output,
@@ -529,12 +552,14 @@ class VectaraToolFactory:
 
         base_params = [
             inspect.Parameter("query", inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=str),
-            inspect.Parameter("top_k", inspect.Parameter.POSITIONAL_OR_KEYWORD, default=10, annotation=int),
+            inspect.Parameter("top_k", inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=int),
             inspect.Parameter("summarize", inspect.Parameter.POSITIONAL_OR_KEYWORD, default=True, annotation=bool),
         ]
         search_tool_extra_desc = tool_description + "\n" + """
+        This tool is meant to perform a search for relevant documents, it is not meant for asking questions.
         The response includes metadata about each relevant document.
-        If summarize=True, it also includes a summary of each document.
+        If summarize=True, it also includes a summary of each document, but takes a lot longer to respond,
+        so avoid using it unless necessary.
         """
 
         tool = _create_tool_from_dynamic_function(
@@ -905,7 +930,7 @@ class ToolsFactory:
         user: str = "postgres",
         password: str = "Password",
         dbname: str = "postgres",
-        max_rows: int = 500,
+        max_rows: int = 1000,
     ) -> List[VectaraTool]:
         """
         Returns a list of database tools.
@@ -923,24 +948,16 @@ class ToolsFactory:
             dbname (str, optional): The database name. Defaults to "postgres".
                You must specify either the sql_database object or the scheme, host, port, user, password, and dbname.
             max_rows (int, optional): if specified, instructs the load_data tool to never return more than max_rows
-               rows. Defaults to 500.
+               rows. Defaults to 1000.
 
         Returns:
             List[VectaraTool]: A list of VectaraTool objects.
         """
         if sql_database:
-            tools = self.get_llama_index_tools(
-                tool_package_name="database",
-                tool_spec_name="DatabaseToolSpec",
-                tool_name_prefix=tool_name_prefix,
-                sql_database=sql_database,
-            )
+            dbt = DatabaseTools(sql_database=sql_database)
         else:
             if scheme in ["postgresql", "mysql", "sqlite", "mssql", "oracle"]:
-                tools = self.get_llama_index_tools(
-                    tool_package_name="database",
-                    tool_spec_name="DatabaseToolSpec",
-                    tool_name_prefix=tool_name_prefix,
+                dbt = DatabaseTools(
                     scheme=scheme,
                     host=host,
                     port=port,
@@ -955,28 +972,19 @@ class ToolsFactory:
                 )
 
         # Update tools with description
+        tools = dbt.to_tool_list()
+        vtools = []
         for tool in tools:
             if content_description:
                 tool.metadata.description = (
                     tool.metadata.description + f"The database tables include data about {content_description}."
                 )
-
-        # Add two new tools: load_sample_data and load_unique_values
-        load_data_tool_index = next(i for i, t in enumerate(tools) if t.metadata.name.endswith("load_data"))
-        load_data_fn_original = tools[load_data_tool_index].fn
-
-        load_data_fn = DBLoadData(load_data_fn_original, max_rows=max_rows)
-        load_data_fn.__name__ = f"{tool_name_prefix}_load_data"
-        load_data_tool = self.create_tool(load_data_fn, ToolType.QUERY)
-
-        sample_data_fn = DBLoadSampleData(load_data_fn_original)
-        sample_data_fn.__name__ = f"{tool_name_prefix}_load_sample_data"
-        sample_data_tool = self.create_tool(sample_data_fn, ToolType.QUERY)
-
-        load_unique_values_fn = DBLoadUniqueValues(load_data_fn_original)
-        load_unique_values_fn.__name__ = f"{tool_name_prefix}_load_unique_values"
-        load_unique_values_tool = self.create_tool(load_unique_values_fn, ToolType.QUERY)
-
-        tools[load_data_tool_index] = load_data_tool
-        tools.extend([sample_data_tool, load_unique_values_tool])
-        return tools
+            if len(tool_name_prefix) > 0:
+                tool.metadata.name = tool_name_prefix + "_" + tool.metadata.name
+            vtool = VectaraTool(
+                tool_type=ToolType.QUERY,
+                fn=tool.fn, async_fn=tool.async_fn,
+                metadata=tool.metadata
+            )
+            vtools.append(vtool)
+        return vtools

vectara_agentic/utils.py

@@ -6,9 +6,9 @@ from typing import Tuple, Callable, Optional
 from functools import lru_cache
 from inspect import signature
 import json
-import requests
-
+import asyncio
 import tiktoken
+import aiohttp
 
 from llama_index.core.llms import LLM
 from llama_index.llms.openai import OpenAI
@@ -101,13 +101,16 @@ def get_llm(
                      max_tokens=max_tokens
             )
     elif model_provider == ModelProvider.ANTHROPIC:
-        llm = Anthropic(model=model_name, temperature=0, max_tokens=max_tokens)
+        llm = Anthropic(
+            model=model_name, temperature=0,
+            max_tokens=max_tokens, cache_idx=2,
+        )
     elif model_provider == ModelProvider.GEMINI:
         from llama_index.llms.gemini import Gemini
         llm = Gemini(
             model=model_name, temperature=0,
             is_function_calling_model=True,
-            max_tokens=max_tokens
+            max_tokens=max_tokens,
         )
     elif model_provider == ModelProvider.TOGETHER:
         from llama_index.llms.together import TogetherLLM
@@ -159,7 +162,7 @@ def remove_self_from_signature(func):
     func.__signature__ = new_sig
     return func
 
-def summarize_vectara_document(corpus_key: str, vectara_api_key, doc_id: str) -> str:
+async def summarize_vectara_document(corpus_key: str, vectara_api_key: str, doc_id: str) -> str:
     """
     Summarize a document in a Vectara corpus using the Vectara API.
     """
@@ -175,8 +178,32 @@ def summarize_vectara_document(corpus_key: str, vectara_api_key, doc_id: str) ->
         'Accept': 'application/json',
         'x-api-key': vectara_api_key
     }
-
-    response = requests.request("POST", url, headers=headers, data=payload, timeout=60)
-    if response.status_code != 200:
-        return f"Vectara Summarization failed with error code {response.status_code}, error={response.text}"
+    timeout = aiohttp.ClientTimeout(total=60)
+    async with aiohttp.ClientSession(timeout=timeout) as session:
+        async with session.post(url, headers=headers, data=payload) as response:
+            if response.status != 200:
+                error_json = await response.json()
+                return (
+                    f"Vectara Summarization failed with error code {response.status}, "
+                    f"error={error_json['messages'][0]}"
+                )
+            data = await response.json()
+            return data["summary"]
     return json.loads(response.text)["summary"]
+
+async def summarize_documents(
+    vectara_corpus_key: str,
+    vectara_api_key: str,
+    doc_ids: list[str]
+) -> dict[str, str]:
+    """
+    Summarize multiple documents in a Vectara corpus using the Vectara API.
+    """
+    if not doc_ids:
+        return {}
+    tasks = [
+        summarize_vectara_document(vectara_corpus_key, vectara_api_key, doc_id)
+        for doc_id in doc_ids
+    ]
+    summaries = await asyncio.gather(*tasks, return_exceptions=True)
+    return dict(zip(doc_ids, summaries))

{vectara_agentic-0.2.7.dist-info → vectara_agentic-0.2.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vectara_agentic
-Version: 0.2.7
+Version: 0.2.8
 Summary: A Python package for creating AI Assistants and AI Agents with Vectara
 Home-page: https://github.com/vectara/py-vectara-agentic
 Author: Ofer Mendelevitch
@@ -16,7 +16,7 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: llama-index==0.12.25
+Requires-Dist: llama-index==0.12.26
 Requires-Dist: llama-index-indices-managed-vectara==0.4.2
 Requires-Dist: llama-index-agent-llm-compiler==0.3.0
 Requires-Dist: llama-index-agent-lats==0.3.0

{vectara_agentic-0.2.7.dist-info → vectara_agentic-0.2.8.dist-info}/RECORD

@@ -1,28 +1,29 @@
 tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/endpoint.py,sha256=frnpdZQpnuQNNKNYgAn2rFTarNG8MCJaNA77Bw_W22A,1420
-tests/test_agent.py,sha256=CU7Zdb1J5kFjrkkIeDr70W1wsikYuOvoYKGYpg4EPxM,6678
+tests/test_agent.py,sha256=t4omKBg9207hpT8b05v9TwuXJCM4knYSTdsXe740eho,4845
 tests/test_agent_planning.py,sha256=_mj73TNP9yUjkUJ-X31r-cQYreJ4qatXOtMrRvVpF4Y,2411
 tests/test_agent_type.py,sha256=JM0Q2GBGHSADoBacz_DW551zWSfbpf7qa8xXqtyWsc4,5671
 tests/test_fallback.py,sha256=M5YD7NHZ0joVU1frYIr9_OiRAIje5mrXrYVcekzlyGs,2829
 tests/test_private_llm.py,sha256=CY-_rCpxGUuxnZ3ypkodw5Jj-sJCNdh6rLbCvULwuJI,2247
+tests/test_serialization.py,sha256=Ed23GN2zhSJNdPFrVK4aqLkOhJKviczR_o0t-r9TuRI,4762
 tests/test_tools.py,sha256=IVKn0HoS2erTCr1mOEGzTkktiY0PCfKNvqnD_pizjOg,3977
 tests/test_workflow.py,sha256=lVyrVHdRO5leYNbYtHTmKqMX0c8_xehCpUA7cXQKVsc,2175
 vectara_agentic/__init__.py,sha256=2GLDS3U6KckK-dBRl9v_x1kSV507gEhjOfuMmmu0Qxg,850
 vectara_agentic/_callback.py,sha256=5PfqjLmuaZIR6dnqmhniTD_zwCgfi7kOu-nexb6Kss4,9688
 vectara_agentic/_observability.py,sha256=fTL3KW0jQU-_JSpFgjO6-XzgDut_oiq9kt4QR-FkSqU,3804
 vectara_agentic/_prompts.py,sha256=LYyiOAiC8imz3U7MSJiuCYAP39afsp7ycXY7-9biyJI,9314
-vectara_agentic/_version.py,sha256=FGUM5lA5uZpmWWB52dt2AMCqWcU0M9b-2BB-raX-EN4,65
-vectara_agentic/agent.py,sha256=nbBl66n56kjEZX4Zconb9IZjESzpjBZIEQdL4uLfurI,43333
+vectara_agentic/_version.py,sha256=HOBvs3gmojKxd7sNMHt6Q-0_rlFpgzlI1gXNZOS_Fqc,65
+vectara_agentic/agent.py,sha256=ioC6EN86_d7SS1jEZ6CUe6OtetuGmLdWftj5bklPfMs,43522
 vectara_agentic/agent_config.py,sha256=y1hSvU5ns0cE2R7BqF65LFstixF1ytJcoVgicGXo7w0,3691
 vectara_agentic/agent_endpoint.py,sha256=QIMejCLlpW2qzXxeDAxv3anF46XMDdVMdKGWhJh3azY,1996
-vectara_agentic/db_tools.py,sha256=VUdcjDFPwauFd2A92mXNYZnCjeMiTzcTka7S5At_3oQ,3595
-vectara_agentic/sub_query_workflow.py,sha256=eS1S7l5PdyLPLZqxUJSR0oM2kvHb4raPGHk8t8td9sc,10939
-vectara_agentic/tools.py,sha256=RpPGWiPHe-9ZiOxNz389W-gNWxegg7m4RlEx_pH9_W0,42881
+vectara_agentic/db_tools.py,sha256=Go03bzma9m-qDH0CPP8hWhf1nu_4S6s7ke0jGqz58Pk,10296
+vectara_agentic/sub_query_workflow.py,sha256=3WoVnryR2NXyYXbLDM1XVLd7DtbCG0jgrVqeDUN4YNQ,10943
+vectara_agentic/tools.py,sha256=Mm2qfJZWnbNa9G-ycYMP7NPLSo4uUJ9_y45YmXxtlSc,42571
 vectara_agentic/tools_catalog.py,sha256=oiw3wAfbpFhh0_6rMvZsyPqWV6QIzHqhZCNzqRxuyV8,4818
 vectara_agentic/types.py,sha256=HcS7vR8P2v2xQTlOc6ZFV2vvlr3OpzSNWhtcLMxqUZc,1792
-vectara_agentic/utils.py,sha256=U4VWCyrvpXfPb9SJpd4Xj7rJCN-cZCNReNm9_uQjnlk,6759
-vectara_agentic-0.2.7.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-vectara_agentic-0.2.7.dist-info/METADATA,sha256=z-IFDKlGmNh9QSIUV4xOOCYKF-PrTUnDukp8M5BNMe4,25046
-vectara_agentic-0.2.7.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-vectara_agentic-0.2.7.dist-info/top_level.txt,sha256=Y7TQTFdOYGYodQRltUGRieZKIYuzeZj2kHqAUpfCUfg,22
-vectara_agentic-0.2.7.dist-info/RECORD,,
+vectara_agentic/utils.py,sha256=nBQqVb4_UNummqVz28DHm3VaKzy8OAq-xSjhU23uxWU,7646
+vectara_agentic-0.2.8.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+vectara_agentic-0.2.8.dist-info/METADATA,sha256=IV5fm77XOPOvqfcpCZUKRxq9QgnoF3mPu-om_sTKEK8,25046
+vectara_agentic-0.2.8.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+vectara_agentic-0.2.8.dist-info/top_level.txt,sha256=Y7TQTFdOYGYodQRltUGRieZKIYuzeZj2kHqAUpfCUfg,22
+vectara_agentic-0.2.8.dist-info/RECORD,,