|
from typing import Any, Optional |
|
|
|
from omagent_core.utils.registry import registry |
|
from pydantic import Field |
|
from pymilvus import MilvusClient |
|
|
|
from .base import ConnectorBase |
|
|
|
|
|
@registry.register_connector() |
|
class MilvusConnector(ConnectorBase): |
|
host: str = Field(default="./db.db") |
|
port: int = Field(default=19530) |
|
password: str = Field(default="") |
|
username: Optional[str] = Field(default="default") |
|
db: Optional[str] = Field(default="default") |
|
alias: Optional[str] = Field(default="alias") |
|
|
|
def model_post_init(self, __context: Any) -> None: |
|
try: |
|
self._client = MilvusClient( |
|
uri=self.host, |
|
user=self.username, |
|
password=self.password, |
|
db_name=self.db, |
|
) |
|
except Exception as e: |
|
raise ConnectionError( |
|
f"Connection to Milvus failed. Please check your connector config in container.yaml. \n Error Message: {e}" |
|
) |
|
|
|
def check_connection(self) -> bool: |
|
"""Check if the connection to Milvus is valid.""" |
|
|
|
self._client.list_collections() |
|
|
