Weaviate Agents extras integration

requirements-devel.txt

@@ -48,4 +48,6 @@ flake8
 flake8-bugbear==24.12.12
 flake8-comprehensions==3.16.0
 flake8-builtins==2.5.0
-black==24.10.0
+black==24.10.0
+
+weaviate-agents>=0.2.0,<1.0.0

setup.cfg

@@ -40,6 +40,9 @@ install_requires =
     grpcio-health-checking>=1.66.2,<2.0.0
 python_requires = >=3.9
 
+[options.extras_require]
+agents =
+    weaviate-agents >=0.2.5, <1.0.0
 
 [options.package_data]
 # If any package or subpackage contains *.txt, *.rst or *.md files, include them:

weaviate/__init__.py

@@ -4,27 +4,14 @@
 
 import os
 import sys
-from importlib.metadata import version, PackageNotFoundError
+from importlib.metadata import PackageNotFoundError, version
 from typing import Any
 
 try:
     __version__ = version("weaviate-client")
 except PackageNotFoundError:
     __version__ = "unknown version"
 
-from .client import Client, WeaviateAsyncClient, WeaviateClient
-from .collections.batch.client import BatchClient, ClientBatchingContextManager
-from .connect.helpers import (
-    connect_to_custom,
-    connect_to_embedded,
-    connect_to_local,
-    connect_to_wcs,
-    connect_to_weaviate_cloud,
-    use_async_with_custom,
-    use_async_with_embedded,
-    use_async_with_local,
-    use_async_with_weaviate_cloud,
-)
 from . import (
     auth,
     backup,
@@ -38,6 +25,19 @@
     outputs,
     types,
 )
+from .client import Client, WeaviateAsyncClient, WeaviateClient
+from .collections.batch.client import BatchClient, ClientBatchingContextManager
+from .connect.helpers import (
+    connect_to_custom,
+    connect_to_embedded,
+    connect_to_local,
+    connect_to_wcs,
+    connect_to_weaviate_cloud,
+    use_async_with_custom,
+    use_async_with_embedded,
+    use_async_with_local,
+    use_async_with_weaviate_cloud,
+)
 
 if not sys.warnoptions:
     from warnings import simplefilter
@@ -76,6 +76,15 @@
     "use_async_with_weaviate_cloud",
 ]
 
+try:
+    import weaviate_agents as agents
+
+    sys.modules["weaviate.agents"] = agents
+    __all__.append("agents")
+except ImportError:
+    pass
+
+
 deprs = [
     "Collection",
     "AuthClientCredentials",

weaviate/agents/__init__.py

@@ -0,0 +1,6 @@
+from weaviate.exceptions import WeaviateAgentsNotInstalledError
+
+try:
+    from weaviate_agents import *  # type: ignore # noqa: F403, F401
+except ImportError:
+    raise WeaviateAgentsNotInstalledError

weaviate/agents/__init__.pyi

@@ -0,0 +1,6 @@
+from weaviate.exceptions import WeaviateAgentsNotInstalledError
+
+try:
+    from weaviate_agents import *
+except ImportError:
+    raise WeaviateAgentsNotInstalledError

weaviate/agents/base.pyi

@@ -0,0 +1,6 @@
+from weaviate.exceptions import WeaviateAgentsNotInstalledError
+
+try:
+    from weaviate_agents.base import *
+except ImportError:
+    raise WeaviateAgentsNotInstalledError

weaviate/agents/classes/__init__.pyi

@@ -0,0 +1,6 @@
+from weaviate.exceptions import WeaviateAgentsNotInstalledError
+
+try:
+    from weaviate_agents.classes import *
+except ImportError:
+    raise WeaviateAgentsNotInstalledError

weaviate/agents/errors.pyi

@@ -0,0 +1,6 @@
+from weaviate.exceptions import WeaviateAgentsNotInstalledError
+
+try:
+    from weaviate_agents.errors import *
+except ImportError:
+    raise WeaviateAgentsNotInstalledError

weaviate/agents/personalization/__init__.pyi

@@ -0,0 +1,6 @@
+from weaviate.exceptions import WeaviateAgentsNotInstalledError
+
+try:
+    from weaviate_agents.personalization import *  # type: ignore
+except ImportError:
+    raise WeaviateAgentsNotInstalledError

weaviate/agents/query/__init__.pyi

@@ -0,0 +1,6 @@
+from weaviate.exceptions import WeaviateAgentsNotInstalledError
+
+try:
+    from weaviate_agents.query import *
+except ImportError:
+    raise WeaviateAgentsNotInstalledError

weaviate/agents/transformation/__init__.pyi

@@ -0,0 +1,6 @@
+from weaviate.exceptions import WeaviateAgentsNotInstalledError
+
+try:
+    from weaviate_agents.transformation import *
+except ImportError:
+    raise WeaviateAgentsNotInstalledError

weaviate/agents/utils.pyi

@@ -0,0 +1,6 @@
+from weaviate.exceptions import WeaviateAgentsNotInstalledError
+
+try:
+    from weaviate_agents.utils import *  # type: ignore
+except ImportError:
+    raise WeaviateAgentsNotInstalledError

weaviate/exceptions.py

@@ -3,7 +3,7 @@
 """
 
 from json.decoder import JSONDecodeError
-from typing import Union, Tuple
+from typing import Tuple, Union
 
 import httpx
 from grpc.aio import AioRpcError  # type: ignore
@@ -371,3 +371,12 @@ class InsufficientPermissionsError(UnexpectedStatusCodeError):
 
     def __init__(self, res: Union[httpx.Response, AioRpcError]) -> None:
         super().__init__("forbidden", res)
+
+
+class WeaviateAgentsNotInstalledError(ImportError):
+    """Error raised when trying to use Weaviate Agents without the required dependencies."""
+
+    def __init__(self) -> None:
+        super().__init__(
+            'Weaviate Agents (Alpha) functionality requires additional dependencies. Please install them using: "pip install weaviate-client[agents]"'
+        )