iatoolkit 0.7.4__py3-none-any.whl → 0.7.6__py3-none-any.whl

repositories/models.py

@@ -0,0 +1,309 @@
+# Copyright (c) 2024 Fernando Libedinsky
+# Product: IAToolkit
+#
+# IAToolkit is open source software.
+
+from sqlalchemy import Column, Integer, String, DateTime, Enum, Text, JSON, Boolean, ForeignKey, Table
+from sqlalchemy.orm import DeclarativeBase
+from sqlalchemy.orm import relationship, class_mapper, declarative_base
+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)
+
+    # encrypted api-key
+    openai_api_key = Column(String, nullable=True)
+    gemini_api_key = Column(String, nullable=True)
+
+    logo_file = Column(String(128), nullable=True, default='')
+    parameters = Column(JSON, nullable=True, default={})
+    created_at = Column(DateTime, default=datetime.now)
+    allow_jwt = Column(Boolean, default=False, 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)
+    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)
+    local_user_id = Column(Integer, default=0, nullable=True)
+    external_user_id = 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")

repositories/profile_repo.py

@@ -0,0 +1,118 @@
+# Copyright (c) 2024 Fernando Libedinsky
+# Product: IAToolkit
+#
+# IAToolkit is open source software.
+
+from repositories.models import User, Company, ApiKey, UserFeedback
+from injector import inject
+from 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:
+            company.parameters = new_company.parameters
+            company.logo_file = new_company.logo_file
+        else:
+            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()
+
+
+
+

repositories/tasks_repo.py

@@ -0,0 +1,52 @@
+# Copyright (c) 2024 Fernando Libedinsky
+# Product: IAToolkit
+#
+# IAToolkit is open source software.
+
+from injector import inject
+from datetime import datetime
+from repositories.models import Task, TaskStatus, TaskType
+from repositories.database_manager import DatabaseManager
+from sqlalchemy import or_
+
+
+class TaskRepo:
+    @inject
+    def __init__(self, db_manager: DatabaseManager):
+        self.session = db_manager.get_session()
+
+    def create_task(self, new_task: Task) -> Task:
+        self.session.add(new_task)
+        self.session.commit()
+        return new_task
+
+    def update_task(self, task: Task) -> Task:
+        self.session.commit()
+        return task
+
+    def get_task_by_id(self, task_id: int):
+        return self.session.query(Task).filter_by(id=task_id).first()
+
+    def create_or_update_task_type(self, new_task_type: TaskType):
+        task_type = self.session.query(TaskType).filter_by(name=new_task_type.name).first()
+        if task_type:
+            task_type.prompt_template = new_task_type.prompt_template
+            task_type.template_args = new_task_type.template_args
+        else:
+            self.session.add(new_task_type)
+            task_type = new_task_type
+        self.session.commit()
+        return task_type
+
+    def get_task_type(self, name: str):
+        task_type = self.session.query(TaskType).filter_by(name=name).first()
+        return task_type
+
+    def get_pending_tasks(self, company_id: int):
+        now = datetime.now()
+        tasks = self.session.query(Task).filter(
+            Task.company_id == company_id,
+            Task.status == TaskStatus.pendiente,
+            or_(Task.execute_at == None, Task.execute_at <= now)
+        ).all()
+        return tasks

repositories/vs_repo.py

@@ -0,0 +1,139 @@
+# Copyright (c) 2024 Fernando Libedinsky
+# Product: IAToolkit
+#
+# IAToolkit is open source software.
+
+from sqlalchemy import  text
+from huggingface_hub import InferenceClient
+from injector import inject
+from common.exceptions import IAToolkitException
+from repositories.database_manager import DatabaseManager
+from repositories.models import Document, VSDoc
+import os
+import logging
+
+class VSRepo:
+    @inject
+    def __init__(self, db_manager: DatabaseManager):
+        self.session = db_manager.get_session()
+
+        # Inicializar el modelo de embeddings
+        self.embedder = InferenceClient(
+            model="sentence-transformers/all-MiniLM-L6-v2",
+            token=os.getenv('HF_TOKEN'))
+
+
+    def add_document(self, vs_chunk_list: list[VSDoc]):
+        try:
+            for doc in vs_chunk_list:
+                # calculate the embedding for the text
+                doc.embedding = self.embedder.feature_extraction(doc.text)
+                self.session.add(doc)
+            self.session.commit()
+        except Exception as e:
+            logging.error(f"Error insertando documentos en PostgreSQL: {str(e)}")
+            self.session.rollback()
+            raise IAToolkitException(IAToolkitException.ErrorType.VECTOR_STORE_ERROR,
+                               f"Error insertando documentos en PostgreSQL: {str(e)}")
+
+    def query(self,
+              company_id: int,
+              query_text: str,
+              n_results=5,
+              metadata_filter=None
+              ) -> list[Document]:
+        """
+        search documents similar to the query for a company
+
+        Args:
+            company_id:
+            query_text: query text
+            n_results: max number of results to return
+            metadata_filter:  (ej: {"document_type": "certificate"})
+
+        Returns:
+            list of documents matching the query and filters
+        """
+        # Generate the embedding with the query text
+        query_embedding = self.embedder.feature_extraction([query_text])[0]
+
+        try:
+            # build the SQL query
+            sql_query_parts = ["""
+                               SELECT iat_documents.id, \
+                                      iat_documents.filename, \
+                                      iat_documents.content, \
+                                      iat_documents.content_b64, \
+                                      iat_documents.meta
+                               FROM iat_vsdocs, \
+                                    iat_documents
+                               WHERE iat_vsdocs.company_id = :company_id
+                                 AND iat_vsdocs.document_id = iat_documents.id \
+                               """]
+
+            # query parameters
+            params = {
+                "company_id": company_id,
+                "query_embedding": query_embedding,
+                "n_results": n_results
+            }
+
+            # add metadata filter, if exists
+            if metadata_filter and isinstance(metadata_filter, dict):
+                for key, value in metadata_filter.items():
+                    # Usar el operador ->> para extraer el valor del JSON como texto.
+                    # La clave del JSON se interpola directamente.
+                    # El valor se pasa como parámetro para evitar inyección SQL.
+                    param_name = f"value_{key}_filter"
+                    sql_query_parts.append(f" AND documents.meta->>'{key}' = :{param_name}")
+                    params[param_name] = str(value)     # parametros como string
+
+            # join all the query parts
+            sql_query = "".join(sql_query_parts)
+
+            # add sorting and limit of results
+            sql_query += " ORDER BY embedding <-> :query_embedding LIMIT :n_results"
+
+            logging.debug(f"Executing SQL query: {sql_query}")
+            logging.debug(f"With parameters: {params}")
+
+            # execute the query
+            result = self.session.execute(text(sql_query), params)
+
+            rows = result.fetchall()
+            vs_documents = []
+
+            for row in rows:
+                # create the document object with the data
+                meta_data = row[4] if len(row) > 4 and row[4] is not None else {}
+                doc = Document(
+                    id=row[0],
+                    company_id=company_id,
+                    filename=row[1],
+                    content=row[2],
+                    content_b64=row[3],
+                    meta=meta_data
+                )
+                vs_documents.append(doc)
+
+            return self.remove_duplicates_by_id(vs_documents)
+
+        except Exception as e:
+            logging.error(f"Error en la consulta de documentos: {str(e)}")
+            logging.error(f"Failed SQL: {sql_query}")
+            logging.error(f"Failed params: {params}")
+            raise IAToolkitException(IAToolkitException.ErrorType.VECTOR_STORE_ERROR,
+                               f"Error en la consulta: {str(e)}")
+        finally:
+            self.session.close()
+
+    def remove_duplicates_by_id(self, objects):
+        unique_by_id = {}
+        result = []
+
+        for obj in objects:
+            if obj.id not in unique_by_id:
+                unique_by_id[obj.id] = True
+                result.append(obj)
+
+        return result

views/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) 2024 Fernando Libedinsky
+# Product: IAToolkit
+#
+# IAToolkit is open source software.
+