btech-agent-lab 1.9.0__py3-none-any.whl

agent_lab/core/container.py

@@ -0,0 +1,220 @@
+import hvac
+from dependency_injector import containers, providers
+
+from agent_lab.domain.repositories.agents import AgentRepository, AgentSettingRepository
+from agent_lab.domain.repositories.attachments import AttachmentRepository
+from agent_lab.domain.repositories.integrations import IntegrationRepository
+from agent_lab.domain.repositories.language_models import (
+    LanguageModelRepository,
+    LanguageModelSettingRepository,
+)
+from agent_lab.domain.repositories.messages import MessageRepository
+from agent_lab.infrastructure.database.checkpoints import GraphPersistenceFactory
+from agent_lab.infrastructure.database.sql import Database
+from agent_lab.infrastructure.database.vectors import DocumentRepository
+from agent_lab.infrastructure.metrics.tracer import (
+    Tracer,
+    service_name,
+    service_version,
+    excluded_paths,
+)
+from agent_lab.infrastructure.metrics.tracing_backends import (
+    LangfuseTracingBackend,
+    LangWatchTracingBackend,
+)
+from agent_lab.interface.mcp.coordinator_planner_supervisor_tool_registrar import (
+    CoordinatorPlannerSupervisorToolRegistrar,
+)
+from agent_lab.interface.mcp.default_tool_registrar import DefaultToolRegistrar
+from agent_lab.interface.mcp.prompt_registry import PromptRegistry
+from agent_lab.interface.mcp.user_prompt_resolver import UserPromptResolver
+from agent_lab.services.agent_settings import AgentSettingService
+from agent_lab.services.agent_types.base import AgentUtils
+from agent_lab.services.agent_types.registry import AgentRegistry
+from agent_lab.services.agents import AgentService
+from agent_lab.services.attachments import AttachmentService
+from agent_lab.services.auth import AuthService
+from agent_lab.services.integrations import IntegrationService
+from agent_lab.services.language_model_settings import LanguageModelSettingService
+from agent_lab.services.language_models import LanguageModelService
+from agent_lab.services.messages import MessageService
+from agent_lab.services.tasks import TaskNotificationService
+
+
+class Container(containers.DeclarativeContainer):
+    wiring_config = containers.WiringConfiguration(
+        modules=[
+            "agent_lab.interface.api.agents.endpoints",
+            "agent_lab.interface.api.attachments.endpoints",
+            "agent_lab.interface.api.auth.endpoints",
+            "agent_lab.interface.api.integrations.endpoints",
+            "agent_lab.interface.api.language_models.endpoints",
+            "agent_lab.interface.api.messages.endpoints",
+        ]
+    )
+
+    # Populated at bootstrap by create_app() via load_config(); declaring the
+    # container performs no I/O so the library is import-safe.
+    config = providers.Configuration()
+
+    db = providers.Singleton(Database, db_url=config.db.url)
+
+    graph_persistence_factory = providers.Singleton(
+        GraphPersistenceFactory, db_checkpoints=config.db.checkpoints
+    )
+
+    vault_client = providers.Singleton(
+        hvac.Client, url=config.vault.url, token=config.vault.token, verify=False
+    )
+
+    document_repository = providers.Factory(
+        DocumentRepository, db_url=config.db.vectors
+    )
+
+    auth_service = providers.Factory(
+        AuthService,
+        enabled=config.auth.enabled,
+        url=config.auth.url,
+        realm=config.auth.realm,
+        client_id=config.auth.client_id,
+        client_secret=config.auth.client_secret,
+    )
+
+    integration_repository = providers.Factory(
+        IntegrationRepository,
+        db=db,
+        vault_client=vault_client,
+    )
+
+    integration_service = providers.Factory(
+        IntegrationService,
+        integration_repository=integration_repository,
+    )
+
+    task_notification_service = providers.Factory(
+        TaskNotificationService,
+        redis_url=config.broker.url,
+    )
+
+    language_model_setting_repository = providers.Factory(
+        LanguageModelSettingRepository,
+        db=db,
+    )
+
+    language_model_setting_service = providers.Factory(
+        LanguageModelSettingService,
+        language_model_setting_repository=language_model_setting_repository,
+    )
+
+    language_model_repository = providers.Factory(LanguageModelRepository, db=db)
+
+    language_model_service = providers.Factory(
+        LanguageModelService,
+        language_model_repository=language_model_repository,
+        language_model_setting_service=language_model_setting_service,
+        integration_service=integration_service,
+    )
+
+    attachment_repository = providers.Factory(AttachmentRepository, db=db)
+
+    attachment_service = providers.Factory(
+        AttachmentService,
+        attachment_repository=attachment_repository,
+        document_repository=document_repository,
+        language_model_service=language_model_service,
+        language_model_setting_service=language_model_setting_service,
+        integration_service=integration_service,
+        vault_client=vault_client,
+    )
+
+    agent_setting_repository = providers.Factory(AgentSettingRepository, db=db)
+
+    agent_setting_service = providers.Factory(
+        AgentSettingService,
+        agent_setting_repository=agent_setting_repository,
+    )
+
+    agent_repository = providers.Factory(AgentRepository, db=db)
+
+    agent_service = providers.Factory(
+        AgentService,
+        agent_repository=agent_repository,
+        agent_setting_service=agent_setting_service,
+        language_model_service=language_model_service,
+    )
+
+    message_repository = providers.Factory(MessageRepository, db=db)
+
+    message_service = providers.Factory(
+        MessageService,
+        message_repository=message_repository,
+        agent_service=agent_service,
+        attachment_service=attachment_service,
+    )
+
+    agent_utils = providers.Factory(
+        AgentUtils,
+        config=config,
+        agent_service=agent_service,
+        agent_setting_service=agent_setting_service,
+        attachment_service=attachment_service,
+        language_model_service=language_model_service,
+        language_model_setting_service=language_model_setting_service,
+        integration_service=integration_service,
+        vault_client=vault_client,
+        graph_persistence_factory=graph_persistence_factory,
+        document_repository=document_repository,
+        task_notification_service=task_notification_service,
+    )
+
+    # Agents register themselves via @RegisterAgent and are discovered through
+    # scan-packages/entry points. The registry needs the live container
+    # instance to resolve each agent's extra_deps, so create_app() binds it
+    # via container.agent_registry.override(...) right after instantiation
+    # (providers.Self() resolves to the base class under subclassing, so it
+    # cannot be used here).
+    agent_registry = providers.Singleton(AgentRegistry)
+
+    user_prompt_resolver = providers.Singleton(
+        UserPromptResolver,
+        agent_service=agent_service,
+        agent_setting_service=agent_setting_service,
+    )
+
+    prompt_registry = providers.Singleton(PromptRegistry)
+
+    default_tool_registrar = providers.Singleton(
+        DefaultToolRegistrar,
+        prompt_registry=prompt_registry,
+    )
+
+    coordinator_planner_supervisor_tool_registrar = providers.Singleton(
+        CoordinatorPlannerSupervisorToolRegistrar,
+        user_prompt_resolver=user_prompt_resolver,
+        prompt_registry=prompt_registry,
+    )
+
+    mcp_registrars = providers.List(
+        default_tool_registrar,
+        coordinator_planner_supervisor_tool_registrar,
+    )
+
+    # Tracing frameworks are decoupled behind the TracingBackend abstraction and
+    # composed here as a list (mirroring mcp_registrars). Adding a framework is a
+    # new backend class + one entry below — the Tracer coordinator, the
+    # @trace_agent_message decorator and the agent call sites stay untouched.
+    langfuse_tracing_backend = providers.Singleton(LangfuseTracingBackend)
+
+    langwatch_tracing_backend = providers.Singleton(
+        LangWatchTracingBackend,
+        service_name=service_name,
+        service_version=service_version,
+        excluded_paths=excluded_paths,
+    )
+
+    tracing_backends = providers.List(
+        langfuse_tracing_backend,
+        langwatch_tracing_backend,
+    )
+
+    tracer = providers.Singleton(Tracer, backends=tracing_backends)

agent_lab/domain/exceptions/base.py

@@ -0,0 +1,54 @@
+from fastapi import HTTPException
+
+
+class NotFoundError(HTTPException):
+    entity_name: str
+
+    def __init__(self, entity_id):
+        super().__init__(
+            status_code=404, detail=f"{self.entity_name} not found, id: {entity_id}"
+        )
+
+
+class InvalidFieldError(HTTPException):
+    def __init__(self, field_name, reason):
+        super().__init__(
+            status_code=400, detail=f"Field {field_name} is invalid, reason: {reason}"
+        )
+
+
+class ResourceNotFoundError(HTTPException):
+    def __init__(self, file_path):
+        super().__init__(status_code=500, detail=f"file not found in path {file_path}")
+
+
+class ConfigurationError(HTTPException):
+    def __init__(self, reason):
+        super().__init__(
+            status_code=500, detail=f"Configuration error, reason: {reason}"
+        )
+
+
+class FileToLargeError(HTTPException):
+    def __init__(self, file_size, max_size):
+        super().__init__(
+            status_code=413,
+            detail=f"File size {file_size} exceeds the maximum allowed size of {max_size} bytes",
+        )
+
+
+class FileProcessingError(HTTPException):
+    def __init__(self, file_name, reason):
+        super().__init__(
+            status_code=422, detail=f"Failed to process file {file_name}: {reason}"
+        )
+
+
+class AudioOptimizationError(HTTPException):
+    def __init__(self, reason):
+        super().__init__(status_code=500, detail=f"Audio optimization failed: {reason}")
+
+
+class AuthenticationError(HTTPException):
+    def __init__(self, reason):
+        super().__init__(status_code=401, detail=f"Authentication failed: {reason}")

agent_lab/domain/models.py

@@ -0,0 +1,110 @@
+from sqlalchemy import (
+    TIMESTAMP,
+    Boolean,
+    Column,
+    Enum,
+    ForeignKey,
+    LargeBinary,
+    String,
+    Text,
+    JSON,
+)
+from sqlalchemy.orm import relationship
+
+from agent_lab.infrastructure.database.sql import Base
+
+message_role = Enum("assistant", "human", name="message_role")
+
+
+class Agent(Base):
+    __tablename__ = "agents"
+
+    id = Column(String, primary_key=True)
+    created_at = Column(TIMESTAMP)
+    is_active = Column(Boolean)
+    agent_name = Column(String)
+    agent_type = Column(String)
+    agent_summary = Column(Text)
+    language_model_id = Column(String, ForeignKey("language_models.id"))
+
+    language_model = relationship("LanguageModel", back_populates="agents")
+    settings = relationship("AgentSetting", back_populates="agent")
+    messages = relationship("Message", back_populates="agent")
+
+
+class AgentSetting(Base):
+    __tablename__ = "agent_settings"
+
+    id = Column(String, primary_key=True)
+    agent_id = Column(String, ForeignKey("agents.id"))
+    setting_key = Column(String)
+    setting_value = Column(Text)
+
+    agent = relationship("Agent", back_populates="settings")
+
+
+class Attachment(Base):
+    __tablename__ = "attachments"
+
+    id = Column(String, primary_key=True)
+    created_at = Column(TIMESTAMP)
+    is_active = Column(Boolean)
+    file_name = Column(String)
+    raw_content = Column(LargeBinary)
+    parsed_content = Column(Text)
+    embeddings_collection = Column(String)
+
+    messages = relationship("Message", back_populates="attachment")
+
+
+class LanguageModel(Base):
+    __tablename__ = "language_models"
+
+    id = Column(String, primary_key=True)
+    created_at = Column(TIMESTAMP)
+    is_active = Column(Boolean)
+    language_model_tag = Column(String)
+    integration_id = Column(String, ForeignKey("integrations.id"))
+
+    agents = relationship("Agent", back_populates="language_model")
+    integration = relationship("Integration", back_populates="language_models")
+    settings = relationship("LanguageModelSetting", back_populates="language_model")
+
+
+class LanguageModelSetting(Base):
+    __tablename__ = "language_model_settings"
+
+    id = Column(String, primary_key=True)
+    language_model_id = Column(String, ForeignKey("language_models.id"))
+    setting_key = Column(String)
+    setting_value = Column(String)
+
+    language_model = relationship("LanguageModel", back_populates="settings")
+
+
+class Integration(Base):
+    __tablename__ = "integrations"
+
+    id = Column(String, primary_key=True)
+    created_at = Column(TIMESTAMP)
+    is_active = Column(Boolean)
+    integration_type = Column(String)
+
+    language_models = relationship("LanguageModel", back_populates="integration")
+
+
+class Message(Base):
+    __tablename__ = "messages"
+
+    id = Column(String, primary_key=True)
+    created_at = Column(TIMESTAMP)
+    is_active = Column(Boolean)
+    agent_id = Column(String, ForeignKey("agents.id"))
+    message_role = Column(message_role)
+    message_content = Column(Text)
+    response_data = Column(JSON)
+    attachment_id = Column(String, ForeignKey("attachments.id"))
+    replies_to = Column(String)
+
+    agent = relationship("Agent", back_populates="messages")
+    attachment = relationship("Attachment", back_populates="messages")

agent_lab/domain/repositories/agents.py

@@ -0,0 +1,152 @@
+from datetime import datetime
+from uuid import uuid4
+
+from typing_extensions import Iterator, Optional
+
+from agent_lab.domain.exceptions.base import NotFoundError
+from agent_lab.domain.models import Agent, AgentSetting
+from agent_lab.infrastructure.database.sql import Database
+
+
+class AgentRepository:
+    def __init__(
+        self,
+        db: Database,
+    ) -> None:
+        self.db = db
+
+    def get_all(self, schema: str) -> Iterator[Agent]:
+        with self.db.session(schema_name=schema) as session:
+            return session.query(Agent).filter(Agent.is_active).all()
+
+    def get_by_id(self, agent_id: str, schema: str) -> Agent:
+        with self.db.session(schema_name=schema) as session:
+            agent = (
+                session.query(Agent)
+                .filter(Agent.id == agent_id, Agent.is_active)
+                .first()
+            )
+            if not agent:
+                raise AgentNotFoundError(agent_id)
+            return agent
+
+    def add(
+        self, agent_name: str, agent_type: str, language_model_id: str, schema: str
+    ) -> Agent:
+        gen_id = uuid4()
+        with self.db.session(schema_name=schema) as session:
+            agent = Agent(
+                id=str(gen_id),
+                is_active=True,
+                created_at=datetime.now(),
+                agent_name=agent_name,
+                agent_type=agent_type,
+                agent_summary="",
+                language_model_id=language_model_id,
+            )
+            session.add(agent)
+            session.commit()
+            session.refresh(agent)
+            return agent
+
+    def update_agent(
+        self,
+        agent_id: str,
+        agent_name: str,
+        language_model_id: str,
+        schema: str,
+        agent_summary: Optional[str] = None,
+    ) -> Agent:
+        with self.db.session(schema_name=schema) as session:
+            entity: Agent = (
+                session.query(Agent)
+                .filter(Agent.id == agent_id, Agent.is_active)
+                .first()
+            )
+            if not entity:
+                raise AgentNotFoundError(agent_id)
+
+            entity.agent_name = agent_name
+            entity.language_model_id = language_model_id
+            if agent_summary is not None:
+                entity.agent_summary = agent_summary
+            session.commit()
+            session.refresh(entity)
+            return entity
+
+    def delete_by_id(self, agent_id: str, schema: str) -> None:
+        with self.db.session(schema_name=schema) as session:
+            entity: Agent = (
+                session.query(Agent)
+                .filter(Agent.id == agent_id, Agent.is_active)
+                .first()
+            )
+            if not entity:
+                raise AgentNotFoundError(agent_id)
+
+            entity.is_active = False
+            session.commit()
+
+
+class AgentNotFoundError(NotFoundError):
+    entity_name: str = "Agent"
+
+
+class AgentSettingRepository:
+    def __init__(self, db: Database) -> None:
+        self.db = db
+
+    def get_all(self, agent_id: str, schema: str) -> Iterator[AgentSetting]:
+        with self.db.session(schema_name=schema) as session:
+            return (
+                session.query(AgentSetting)
+                .filter(AgentSetting.agent_id == agent_id)
+                .all()
+            )
+
+    def add(
+        self, agent_id: str, setting_key: str, setting_value: str, schema: str
+    ) -> AgentSetting:
+        gen_id = uuid4()
+        with self.db.session(schema_name=schema) as session:
+            agent_setting = AgentSetting(
+                id=str(gen_id),
+                agent_id=agent_id,
+                setting_key=setting_key,
+                setting_value=setting_value,
+            )
+            session.add(agent_setting)
+            session.commit()
+            session.refresh(agent_setting)
+            return agent_setting
+
+    def update_by_key(
+        self, agent_id: str, setting_key: str, setting_value: str, schema: str
+    ) -> AgentSetting:
+        with self.db.session(schema_name=schema) as session:
+            entity: AgentSetting = (
+                session.query(AgentSetting)
+                .filter(
+                    AgentSetting.agent_id == agent_id,
+                    AgentSetting.setting_key == setting_key,
+                )
+                .first()
+            )
+            if not entity:
+                raise AgentSettingNotFoundError(agent_id)
+
+            entity.setting_value = setting_value
+            session.commit()
+            session.refresh(entity)
+            return entity
+
+    def delete_by_agent_id(self, agent_id: str, schema: str) -> None:
+        with self.db.session(schema_name=schema) as session:
+            session.query(AgentSetting).filter(
+                AgentSetting.agent_id == agent_id
+            ).delete(synchronize_session=False)
+            session.commit()
+
+
+class AgentSettingNotFoundError(NotFoundError):
+    entity_name: str = "AgentSetting"

agent_lab/domain/repositories/attachments.py

@@ -0,0 +1,86 @@
+from datetime import datetime
+from uuid import uuid4
+
+from typing_extensions import Iterator
+
+from agent_lab.domain.exceptions.base import NotFoundError
+from agent_lab.domain.models import Attachment
+from agent_lab.infrastructure.database.sql import Database
+
+
+class AttachmentRepository:
+    def __init__(self, db: Database) -> None:
+        self.db = db
+
+    def get_all(self, schema: str) -> Iterator[Attachment]:
+        with self.db.session(schema_name=schema) as session:
+            return session.query(Attachment).filter(Attachment.is_active).all()
+
+    def get_by_id(self, attachment_id: str, schema: str) -> Attachment:
+        with self.db.session(schema_name=schema) as session:
+            attachment = (
+                session.query(Attachment)
+                .filter(Attachment.id == attachment_id, Attachment.is_active)
+                .first()
+            )
+            if not attachment:
+                raise AttachmentNotFoundError(attachment_id)
+            return attachment
+
+    def add(
+        self,
+        file_name: str,
+        raw_content: bytes,
+        parsed_content: str,
+        schema: str,
+        attachment_id: str = None,
+    ) -> Attachment:
+        if attachment_id is None:
+            attachment_id = str(uuid4())
+
+        with self.db.session(schema_name=schema) as session:
+            attachment = Attachment(
+                id=attachment_id,
+                is_active=True,
+                created_at=datetime.now(),
+                file_name=file_name,
+                raw_content=raw_content,
+                parsed_content=parsed_content,
+            )
+            session.add(attachment)
+            session.commit()
+            session.refresh(attachment)
+            return attachment
+
+    def delete_by_id(self, attachment_id: str, schema: str) -> None:
+        with self.db.session(schema_name=schema) as session:
+            entity: Attachment = (
+                session.query(Attachment)
+                .filter(Attachment.id == attachment_id, Attachment.is_active)
+                .first()
+            )
+            if not entity:
+                raise AttachmentNotFoundError(attachment_id)
+            session.delete(entity)
+            session.commit()
+
+    def update_attachment(
+        self, attachment_id: str, embeddings_collection: str, schema: str
+    ) -> Attachment:
+        with self.db.session(schema_name=schema) as session:
+            entity: Attachment = (
+                session.query(Attachment)
+                .filter(Attachment.id == attachment_id, Attachment.is_active)
+                .first()
+            )
+            if not entity:
+                raise AttachmentNotFoundError(attachment_id)
+
+            entity.embeddings_collection = embeddings_collection
+            session.commit()
+            session.refresh(entity)
+            return entity
+
+
+class AttachmentNotFoundError(NotFoundError):
+    entity_name: str = "Attachment"

agent_lab/domain/repositories/integrations.py

@@ -0,0 +1,72 @@
+from datetime import datetime
+from uuid import uuid4
+
+import hvac
+from typing_extensions import Iterator
+
+from agent_lab.domain.exceptions.base import NotFoundError
+from agent_lab.domain.models import Integration
+from agent_lab.infrastructure.database.sql import Database
+
+
+class IntegrationRepository:
+    def __init__(
+        self,
+        db: Database,
+        vault_client: hvac.Client,
+    ) -> None:
+        self.db = db
+        self.vault_client = vault_client
+
+    def get_all(self, schema: str) -> Iterator[Integration]:
+        with self.db.session(schema_name=schema) as session:
+            return session.query(Integration).filter(Integration.is_active).all()
+
+    def get_by_id(self, integration_id: str, schema: str) -> Integration:
+        with self.db.session(schema_name=schema) as session:
+            integration = (
+                session.query(Integration)
+                .filter(Integration.id == integration_id, Integration.is_active)
+                .first()
+            )
+            if not integration:
+                raise IntegrationNotFoundError(integration_id)
+            return integration
+
+    def add(
+        self, integration_type: str, api_endpoint: str, api_key: str, schema: str
+    ) -> Integration:
+        gen_id = uuid4()
+        self.vault_client.secrets.kv.v2.create_or_update_secret(
+            path=f"integration_{gen_id}",
+            secret={"api_endpoint": api_endpoint, "api_key": api_key},
+        )
+
+        with self.db.session(schema_name=schema) as session:
+            integration = Integration(
+                id=str(gen_id),
+                created_at=datetime.now(),
+                is_active=True,
+                integration_type=integration_type,
+            )
+            session.add(integration)
+            session.commit()
+            session.refresh(integration)
+            return integration
+
+    def delete_by_id(self, integration_id: str, schema: str) -> None:
+        with self.db.session(schema_name=schema) as session:
+            entity: Integration = (
+                session.query(Integration)
+                .filter(Integration.id == integration_id, Integration.is_active)
+                .first()
+            )
+            if not entity:
+                raise IntegrationNotFoundError(integration_id)
+
+            entity.is_active = False
+            session.commit()
+
+
+class IntegrationNotFoundError(NotFoundError):
+    entity_name: str = "Integration"