iatoolkit 0.4.2__py3-none-any.whl → 0.66.2__py3-none-any.whl

iatoolkit/repositories/database_manager.py

@@ -0,0 +1,110 @@
+# Copyright (c) 2024 Fernando Libedinsky
+# Product: IAToolkit
+#
+# IAToolkit is open source software.
+
+# database_manager.py
+from sqlalchemy import create_engine, event, inspect
+from sqlalchemy.orm import sessionmaker, scoped_session
+from sqlalchemy.engine.url import make_url
+from iatoolkit.repositories.models import Base
+from injector import inject
+from pgvector.psycopg2 import register_vector
+
+
+class DatabaseManager:
+    @inject
+    def __init__(self, database_url: str, register_pgvector: bool = True):
+        """
+        Inicializa el gestor de la base de datos.
+        :param database_url: URL de la base de datos.
+        :param echo: Si True, habilita logs de SQL.
+        """
+        self.url = make_url(database_url)
+        if database_url.startswith('sqlite'):       # for tests
+            self._engine = create_engine(database_url, echo=False)
+        else:
+            self._engine = create_engine(
+                database_url,
+                echo=False,
+                pool_size=2,  # per worker
+                max_overflow=3,
+                pool_timeout=30,
+                pool_recycle=1800,
+                pool_pre_ping=True,
+                future=True,
+            )
+        self.SessionFactory = sessionmaker(bind=self._engine,
+                                           autoflush=False,
+                                           autocommit=False,
+                                           expire_on_commit=False)
+        self.scoped_session = scoped_session(self.SessionFactory)
+
+        # REGISTRAR pgvector para cada nueva conexión solo en postgres
+        if register_pgvector and self.url.get_backend_name() == 'postgresql':
+            event.listen(self._engine, 'connect', self.on_connect)
+
+    @staticmethod
+    def on_connect(dbapi_connection, connection_record):
+        """
+        Esta función se ejecuta cada vez que se establece una conexión.
+        dbapi_connection es la conexión psycopg2 real.
+        """
+        register_vector(dbapi_connection)
+
+    def get_session(self):
+        return self.scoped_session()
+
+    def get_connection(self):
+        return self._engine.connect()
+
+    def get_engine(self):
+        return self._engine
+
+    def create_all(self):
+        Base.metadata.create_all(self._engine)
+
+    def drop_all(self):
+        Base.metadata.drop_all(self._engine)
+
+    def remove_session(self):
+        self.scoped_session.remove()
+
+    def get_table_schema(self,
+                         table_name: str,
+                         schema_name: str | None = None,
+                         exclude_columns: list[str] | None = None) -> str:
+        inspector = inspect(self._engine)
+
+        if table_name not in inspector.get_table_names():
+            raise RuntimeError(f"La tabla '{table_name}' no existe en la BD.")
+
+        if exclude_columns is None:
+            exclude_columns = []
+
+        # get all thre table columns
+        columns = inspector.get_columns(table_name)
+
+        # construct a json dictionary with the table definition
+        json_dict = {
+            "table": table_name,
+            "description": f"Definición de la tabla {table_name}.",
+            "fields": []
+        }
+        if schema_name:
+            json_dict["description"] += f"Los detalles de cada campo están en el objeto **`{schema_name}`**."
+
+        # now add every column to the json dictionary
+        for col in columns:
+            name = col["name"]
+
+            # omit the excluded columns.
+            if name in exclude_columns:
+                continue
+
+            json_dict["fields"].append({
+                "name": name,
+                "type": str(col["type"]),
+            })
+
+        return "\n\n" + str(json_dict)

iatoolkit/repositories/document_repo.py

@@ -0,0 +1,33 @@
+# Copyright (c) 2024 Fernando Libedinsky
+# Product: IAToolkit
+#
+# IAToolkit is open source software.
+
+from iatoolkit.repositories.models import Document
+from injector import inject
+from iatoolkit.repositories.database_manager import DatabaseManager
+from iatoolkit.common.exceptions import IAToolkitException
+
+
+class DocumentRepo:
+    @inject
+    def __init__(self, db_manager: DatabaseManager):
+        self.session = db_manager.get_session()
+
+    def insert(self,new_document: Document):
+        self.session.add(new_document)
+        self.session.commit()
+        return new_document
+
+    def get(self, company_id, filename: str ) -> Document:
+        if not company_id or not filename:
+            raise IAToolkitException(IAToolkitException.ErrorType.PARAM_NOT_FILLED,
+                               'Falta empresa o filename')
+
+        return self.session.query(Document).filter_by(company_id=company_id, filename=filename).first()
+
+    def get_by_id(self, document_id: int) -> Document:
+        if not document_id:
+            return None
+
+        return self.session.query(Document).filter_by(id=document_id).first()

iatoolkit/repositories/llm_query_repo.py

@@ -0,0 +1,91 @@
+# Copyright (c) 2024 Fernando Libedinsky
+# Product: IAToolkit
+#
+# IAToolkit is open source software.
+
+from iatoolkit.repositories.models import LLMQuery, Function, Company, Prompt, PromptCategory
+from injector import inject
+from iatoolkit.repositories.database_manager import DatabaseManager
+from sqlalchemy import or_
+
+class LLMQueryRepo:
+    @inject
+    def __init__(self, db_manager: DatabaseManager):
+        self.session = db_manager.get_session()
+
+    def add_query(self, query: LLMQuery):
+        self.session.add(query)
+        self.session.commit()
+        return query
+
+
+    def get_company_functions(self, company: Company) -> list[Function]:
+        return (
+            self.session.query(Function)
+            .filter(
+                Function.is_active.is_(True),
+                or_(
+                    Function.company_id == company.id,
+                    Function.system_function.is_(True)
+                )
+            )
+            .all()
+        )
+
+    def create_or_update_function(self, new_function: Function):
+        function = self.session.query(Function).filter_by(company_id=new_function.company_id,
+                                                 name=new_function.name).first()
+        if function:
+            function.description = new_function.description
+            function.parameters = new_function.parameters
+            function.system_function = new_function.system_function
+        else:
+            self.session.add(new_function)
+            function = new_function
+
+        self.session.commit()
+        return function
+
+    def create_or_update_prompt(self, new_prompt: Prompt):
+        prompt = self.session.query(Prompt).filter_by(company_id=new_prompt.company_id,
+                                                 name=new_prompt.name).first()
+        if prompt:
+            prompt.category_id = new_prompt.category_id
+            prompt.description = new_prompt.description
+            prompt.order = new_prompt.order
+            prompt.active = new_prompt.active
+            prompt.is_system_prompt = new_prompt.is_system_prompt
+            prompt.filename = new_prompt.filename
+            prompt.custom_fields = new_prompt.custom_fields
+        else:
+            self.session.add(new_prompt)
+            prompt = new_prompt
+
+        self.session.commit()
+        return prompt
+
+    def create_or_update_prompt_category(self, new_category: PromptCategory):
+        category = self.session.query(PromptCategory).filter_by(company_id=new_category.company_id,
+                                                      name=new_category.name).first()
+        if category:
+            category.order = new_category.order
+        else:
+            self.session.add(new_category)
+            category = new_category
+
+        self.session.commit()
+        return category
+
+    def get_history(self, company: Company, user_identifier: str) -> list[LLMQuery]:
+        return self.session.query(LLMQuery).filter(
+            LLMQuery.user_identifier == user_identifier,
+        ).filter_by(company_id=company.id).order_by(LLMQuery.created_at.desc()).limit(100).all()
+
+    def get_prompts(self, company: Company) -> list[Prompt]:
+        return self.session.query(Prompt).filter_by(company_id=company.id, is_system_prompt=False).all()
+
+    def get_system_prompts(self) -> list[Prompt]:
+        return self.session.query(Prompt).filter_by(is_system_prompt=True, active=True).order_by(Prompt.order).all()
+
+    def get_prompt_by_name(self, company: Company, prompt_name: str):
+        return self.session.query(Prompt).filter_by(company_id=company.id, name=prompt_name).first()

iatoolkit/repositories/models.py

@@ -0,0 +1,336 @@
+# Copyright (c) 2024 Fernando Libedinsky
+# Product: IAToolkit
+#
+# IAToolkit is open source software.
+
+from sqlalchemy import Column, Integer, BigInteger, String, DateTime, Enum, Text, JSON, Boolean, ForeignKey, Table
+from sqlalchemy.orm import DeclarativeBase
+from sqlalchemy.orm import relationship, class_mapper, declarative_base
+from sqlalchemy.sql import func
+from datetime import datetime
+from pgvector.sqlalchemy import Vector
+from enum import Enum as PyEnum
+import secrets
+import enum
+
+
+# base class for the ORM
+class Base(DeclarativeBase):
+    pass
+
+# relation table for many-to-many relationship between companies and users
+user_company = Table('iat_user_company',
+                     Base.metadata,
+                    Column('user_id', Integer,
+                           ForeignKey('iat_users.id', ondelete='CASCADE'),
+                                primary_key=True),
+                     Column('company_id', Integer,
+                            ForeignKey('iat_companies.id',ondelete='CASCADE'),
+                                primary_key=True),
+                     Column('is_active', Boolean, default=True),
+                     Column('role', String(50), default='user'),  # Para manejar roles por empresa
+                     Column('created_at', DateTime, default=datetime.now)
+                     )
+
+class ApiKey(Base):
+    """Represents an API key for a company to authenticate against the system."""
+    __tablename__ = 'iat_api_keys'
+
+    id = Column(Integer, primary_key=True)
+    company_id = Column(Integer, ForeignKey('iat_companies.id', ondelete='CASCADE'), nullable=False)
+    key = Column(String(128), unique=True, nullable=False, index=True) # La API Key en sí
+    is_active = Column(Boolean, default=True, nullable=False)
+    created_at = Column(DateTime, default=datetime.now)
+    last_used_at = Column(DateTime, nullable=True) # Opcional: para rastrear uso
+
+    company = relationship("Company", back_populates="api_keys")
+
+
+class Company(Base):
+    """Represents a company or tenant in the multi-tenant system."""
+    __tablename__ = 'iat_companies'
+
+    id = Column(Integer, primary_key=True)
+    short_name = Column(String(20), nullable=False, unique=True, index=True)
+    name = Column(String(256), nullable=False)
+    default_language = Column(String(5), nullable=False, default='es')
+
+    # encrypted api-key
+    openai_api_key = Column(String, nullable=True)
+    gemini_api_key = Column(String, nullable=True)
+
+    branding = Column(JSON, nullable=True)
+    onboarding_cards = Column(JSON, nullable=True)
+    parameters = Column(JSON, nullable=True)
+    created_at = Column(DateTime, default=datetime.now)
+    allow_jwt = Column(Boolean, default=True, nullable=True)
+
+    documents = relationship("Document",
+                             back_populates="company",
+                             cascade="all, delete-orphan",
+                             lazy='dynamic')
+    functions = relationship("Function",
+                           back_populates="company",
+                           cascade="all, delete-orphan")
+    vsdocs = relationship("VSDoc",
+                          back_populates="company",
+                          cascade="all, delete-orphan")
+    llm_queries = relationship("LLMQuery",
+                               back_populates="company",
+                               cascade="all, delete-orphan")
+    users = relationship("User",
+                         secondary=user_company,
+                         back_populates="companies")
+    api_keys = relationship("ApiKey",
+                            back_populates="company",
+                            cascade="all, delete-orphan")
+
+    tasks = relationship("Task", back_populates="company")
+    feedbacks = relationship("UserFeedback",
+                               back_populates="company",
+                               cascade="all, delete-orphan")
+    prompts = relationship("Prompt",
+                             back_populates="company",
+                             cascade="all, delete-orphan")
+
+    def to_dict(self):
+        return {column.key: getattr(self, column.key) for column in class_mapper(self.__class__).columns}
+
+# users with rights to use this app
+class User(Base):
+    """Represents an IAToolkit user who can be associated with multiple companies."""
+    __tablename__ = 'iat_users'
+
+    id = Column(Integer, primary_key=True)
+    email = Column(String(80), unique=True, nullable=False)
+    first_name = Column(String(50), nullable=False)
+    last_name = Column(String(50), nullable=False)
+    created_at = Column(DateTime, default=datetime.now)
+    password = Column(String, nullable=False)
+    verified = Column(Boolean, nullable=False, default=False)
+    preferred_language = Column(String(5), nullable=True)
+    verification_url = Column(String, nullable=True)
+    temp_code = Column(String, nullable=True)
+
+    companies = relationship(
+        "Company",
+        secondary=user_company,
+        back_populates="users",
+        cascade="all",
+        passive_deletes=True,
+        lazy='dynamic'
+    )
+
+    def to_dict(self):
+        return {
+            'id': self.id,
+            'email': self.email,
+            'first_name': self.first_name,
+            'last_name': self.last_name,
+            'created_at': str(self.created_at),
+            'verified': self.verified,
+            'companies': [company.to_dict() for company in self.companies]
+        }
+
+class Function(Base):
+    """Represents a custom or system function that the LLM can call (tool)."""
+    __tablename__ = 'iat_functions'
+
+    id = Column(Integer, primary_key=True)
+    company_id = Column(Integer,
+                        ForeignKey('iat_companies.id',ondelete='CASCADE'),
+                        nullable=True)
+    name = Column(String(255), nullable=False)
+    system_function = Column(Boolean, default=False)
+    description = Column(Text, nullable=False)
+    parameters = Column(JSON, nullable=False)
+    is_active = Column(Boolean, default=True)
+    created_at = Column(DateTime, default=datetime.now)
+
+    company = relationship('Company', back_populates='functions')
+
+    def to_dict(self):
+        return {column.key: getattr(self, column.key) for column in class_mapper(self.__class__).columns}
+
+
+class Document(Base):
+    """Represents a file or document uploaded by a company for context."""
+    __tablename__ = 'iat_documents'
+
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    company_id = Column(Integer, ForeignKey('iat_companies.id',
+                    ondelete='CASCADE'), nullable=False)
+    filename = Column(String(256), nullable=False, index=True)
+    content = Column(Text, nullable=False)
+    content_b64 = Column(Text, nullable=False)
+    meta = Column(JSON, nullable=True)
+    created_at = Column(DateTime, default=datetime.now)
+
+    company = relationship("Company", back_populates="documents")
+
+    def to_dict(self):
+        return {column.key: getattr(self, column.key) for column in class_mapper(self.__class__).columns}
+
+
+class LLMQuery(Base):
+    """Logs a query made to the LLM, including input, output, and metadata."""
+    __tablename__ = 'iat_queries'
+
+    id = Column(Integer, primary_key=True)
+    company_id = Column(Integer, ForeignKey('iat_companies.id',
+                            ondelete='CASCADE'), nullable=False)
+    user_identifier = Column(String(128), nullable=False)
+    task_id = Column(Integer, default=0, nullable=True)
+    query = Column(Text, nullable=False)
+    output = Column(Text, nullable=False)
+    response = Column(JSON, nullable=True, default={})
+    valid_response = Column(Boolean, nullable=False, default=False)
+    function_calls = Column(JSON, nullable=True, default={})
+    stats = Column(JSON, default={})
+    answer_time = Column(Integer, default=0)
+    created_at = Column(DateTime, default=datetime.now)
+
+    company = relationship("Company", back_populates="llm_queries")
+    tasks = relationship("Task", back_populates="llm_query")
+
+    def to_dict(self):
+        return {column.key: getattr(self, column.key) for column in class_mapper(self.__class__).columns}
+
+
+class VSDoc(Base):
+    """Stores a text chunk and its corresponding vector embedding for similarity search."""
+    __tablename__ = "iat_vsdocs"
+
+    id = Column(Integer, primary_key=True)
+    company_id = Column(Integer, ForeignKey('iat_companies.id',
+                    ondelete='CASCADE'), nullable=False)
+    document_id = Column(Integer, ForeignKey('iat_documents.id',
+                        ondelete='CASCADE'), nullable=False)
+    text = Column(Text, nullable=False)
+    embedding = Column(Vector(384), nullable=False)  # Ajusta la dimensión si es necesario
+
+    company = relationship("Company", back_populates="vsdocs")
+
+    def to_dict(self):
+        return {column.key: getattr(self, column.key) for column in class_mapper(self.__class__).columns}
+
+class TaskStatus(PyEnum):
+    """Enumeration for the possible statuses of a Task."""
+    pendiente = "pendiente"  # task created and waiting to be executed.
+    ejecutado = "ejecutado"  # the IA algorithm has been executed.
+    aprobada = "aprobada"  # validated and approved by human.
+    rechazada = "rechazada"  # validated and rejected by human.
+    fallida = "fallida"  # error executing the IA algorithm.
+
+class TaskType(Base):
+    """Defines a type of task that can be executed, including its prompt template."""
+    __tablename__ = 'iat_task_types'
+
+    id = Column(Integer, primary_key=True)
+    name = Column(String(100), unique=True, nullable=False)
+    prompt_template = Column(String(100), nullable=True)  # Plantilla de prompt por defecto.
+    template_args = Column(JSON, nullable=True)  # Argumentos/prefijos de configuración para el template.
+
+class Task(Base):
+    """Represents an asynchronous task to be executed by the system, often involving an LLM."""
+    __tablename__ = 'iat_tasks'
+
+    id = Column(Integer, primary_key=True)
+    company_id = Column(Integer, ForeignKey("iat_companies.id"))
+
+    user_id = Column(Integer, nullable=True, default=0)
+    task_type_id = Column(Integer, ForeignKey('iat_task_types.id'), nullable=False)
+    status = Column(Enum(TaskStatus, name="task_status_enum"),
+                    default=TaskStatus.pendiente, nullable=False)
+    client_data = Column(JSON, nullable=True, default={})
+    company_task_id = Column(Integer, nullable=True, default=0)
+    execute_at = Column(DateTime, default=datetime.now, nullable=True)
+    llm_query_id = Column(Integer, ForeignKey('iat_queries.id'), nullable=True)
+    callback_url = Column(String(512), default=None, nullable=True)
+    files = Column(JSON, default=[], nullable=True)
+
+    review_user = Column(String(128), nullable=True, default='')
+    review_date = Column(DateTime, nullable=True)
+    comment = Column(Text, nullable=True)
+    approved = Column(Boolean, nullable=False, default=False)
+
+    created_at = Column(DateTime, default=datetime.now)
+    updated_at = Column(DateTime, default=datetime.now)
+
+    task_type = relationship("TaskType")
+    llm_query = relationship("LLMQuery", back_populates="tasks", uselist=False)
+    company = relationship("Company", back_populates="tasks")
+
+class UserFeedback(Base):
+    """Stores feedback and ratings submitted by users for specific interactions."""
+    __tablename__ = 'iat_feedback'
+
+    id = Column(Integer, primary_key=True)
+    company_id = Column(Integer, ForeignKey('iat_companies.id',
+                                            ondelete='CASCADE'), nullable=False)
+    user_identifier = Column(String(128), default='', nullable=True)
+    message = Column(Text, nullable=False)
+    rating = Column(Integer, nullable=False)
+    created_at = Column(DateTime, default=datetime.now)
+
+    company = relationship("Company", back_populates="feedbacks")
+
+
+class PromptCategory(Base):
+    """Represents a category to group and organize prompts."""
+    __tablename__ = 'iat_prompt_categories'
+    id = Column(Integer, primary_key=True)
+    name = Column(String, nullable=False)
+    order = Column(Integer, nullable=False, default=0)
+    company_id = Column(Integer, ForeignKey('iat_companies.id'), nullable=False)
+
+    prompts = relationship("Prompt", back_populates="category", order_by="Prompt.order")
+
+    def __repr__(self):
+        return f"<PromptCategory(name='{self.name}', order={self.order})>"
+
+
+class Prompt(Base):
+    """Represents a system or user-defined prompt template for the LLM."""
+    __tablename__ = 'iat_prompt'
+
+    id = Column(Integer, primary_key=True)
+    company_id = Column(Integer, ForeignKey('iat_companies.id',
+                                            ondelete='CASCADE'), nullable=True)
+    name = Column(String(64), nullable=False)
+    description = Column(String(256), nullable=False)
+    filename = Column(String(256), nullable=False)
+    active = Column(Boolean, default=True)
+    is_system_prompt = Column(Boolean, default=False)
+    order = Column(Integer, nullable=False, default=0)  # Nuevo campo para el orden
+    category_id = Column(Integer, ForeignKey('iat_prompt_categories.id'), nullable=True)
+    custom_fields = Column(JSON, nullable=False, default=[])
+
+    created_at = Column(DateTime, default=datetime.now)
+
+    company = relationship("Company", back_populates="prompts")
+    category = relationship("PromptCategory", back_populates="prompts")
+
+class AccessLog(Base):
+    # Modelo ORM para registrar cada intento de acceso a la plataforma.
+    __tablename__ = 'iat_access_log'
+
+    id = Column(BigInteger, primary_key=True)
+
+    timestamp = Column(DateTime(timezone=True), server_default=func.now(), nullable=False, index=True)
+    company_short_name = Column(String(100), nullable=False, index=True)
+    user_identifier = Column(String(255), index=True)
+
+    # Cómo y el Resultado
+    auth_type = Column(String(20), nullable=False) # 'local', 'external_api', 'redeem_token', etc.
+    outcome = Column(String(10), nullable=False)   # 'success' o 'failure'
+    reason_code = Column(String(50))               # Causa de fallo, ej: 'INVALID_CREDENTIALS'
+
+    # Contexto de la Petición
+    source_ip = Column(String(45), nullable=False)
+    user_agent_hash = Column(String(16))           # Hash corto del User-Agent
+    request_path = Column(String(255), nullable=False)
+
+    def __repr__(self):
+        return (f"<AccessLog(id={self.id}, company='{self.company_short_name}', "
+                f"user='{self.user_identifier}', outcome='{self.outcome}')>")

iatoolkit/repositories/profile_repo.py

@@ -0,0 +1,123 @@
+# Copyright (c) 2024 Fernando Libedinsky
+# Product: IAToolkit
+#
+# IAToolkit is open source software.
+
+from iatoolkit.repositories.models import User, Company, ApiKey, UserFeedback
+from injector import inject
+from iatoolkit.repositories.database_manager import DatabaseManager
+from sqlalchemy.orm import joinedload # Para cargar la relación eficientemente
+
+
+class ProfileRepo:
+    @inject
+    def __init__(self, db_manager: DatabaseManager):
+        self.session = db_manager.get_session()
+
+    def get_user_by_id(self, user_id: int) -> User:
+        user = self.session.query(User).filter_by(id=user_id).first()
+        return user
+
+    def get_user_by_email(self, email: str) -> User:
+        user = self.session.query(User).filter_by(email=email).first()
+        return user
+
+    def create_user(self, new_user: User):
+        self.session.add(new_user)
+        self.session.commit()
+        return new_user
+
+    def save_user(self,existing_user: User):
+        self.session.add(existing_user)
+        self.session.commit()
+        return existing_user
+
+    def update_user(self, email, **kwargs):
+        user = self.session.query(User).filter_by(email=email).first()
+        if not user:
+            return None
+
+        # get the fields for update
+        for key, value in kwargs.items():
+            if hasattr(user, key):  # Asegura que el campo existe en User
+                setattr(user, key, value)
+
+        self.session.commit()
+        return user             # return updated object
+
+    def verify_user(self, email):
+        return self.update_user(email, verified=True)
+
+    def set_temp_code(self, email, temp_code):
+        return self.update_user(email, temp_code=temp_code)
+
+    def reset_temp_code(self, email):
+        return self.update_user(email, temp_code=None)
+
+    def update_password(self, email, hashed_password):
+        return self.update_user(email, password=hashed_password)
+
+    def get_company(self, name: str) -> Company:
+        return self.session.query(Company).filter_by(name=name).first()
+
+    def get_company_by_id(self, company_id: int) -> Company:
+        return self.session.query(Company).filter_by(id=company_id).first()
+
+    def get_company_by_short_name(self, short_name: str) -> Company:
+        return self.session.query(Company).filter(Company.short_name == short_name).first()
+
+    def get_companies(self) -> list[Company]:
+        return self.session.query(Company).all()
+
+    def create_company(self, new_company: Company):
+        company = self.session.query(Company).filter_by(name=new_company.name).first()
+        if company:
+            if company.parameters != new_company.parameters:
+                company.parameters = new_company.parameters
+            if company.branding != new_company.branding:
+                company.branding = new_company.branding
+            if company.onboarding_cards != new_company.onboarding_cards:
+                company.onboarding_cards = new_company.onboarding_cards
+        else:
+            # Si la compañía no existe, la añade a la sesión.
+            self.session.add(new_company)
+            company = new_company
+
+        self.session.commit()
+        return company
+
+    def save_feedback(self, feedback: UserFeedback):
+        self.session.add(feedback)
+        self.session.commit()
+        return feedback
+
+    def create_api_key(self, new_api_key: ApiKey):
+        self.session.add(new_api_key)
+        self.session.commit()
+        return new_api_key
+
+
+    def get_active_api_key_entry(self, api_key_value: str) -> ApiKey | None:
+        """
+        search for an active API Key by its value.
+        returns the entry if found and is active, None otherwise.
+        """
+        try:
+            # Usamos joinedload para cargar la compañía en la misma consulta
+            api_key_entry = self.session.query(ApiKey)\
+                .options(joinedload(ApiKey.company))\
+                .filter(ApiKey.key == api_key_value, ApiKey.is_active == True)\
+                .first()
+            return api_key_entry
+        except Exception:
+            self.session.rollback() # Asegura que la sesión esté limpia tras un error
+            return None
+
+    def get_active_api_key_by_company(self, company: Company) -> ApiKey | None:
+        return self.session.query(ApiKey)\
+                .filter(ApiKey.company == company, ApiKey.is_active == True)\
+                .first()
+
+
+
+