iatoolkit 0.3.9__py3-none-any.whl → 0.107.4__py3-none-any.whl

iatoolkit/repositories/filesystem_asset_repository.py

@@ -0,0 +1,36 @@
+from iatoolkit.common.interfaces.asset_storage import AssetRepository, AssetType
+from pathlib import Path
+
+
+class FileSystemAssetRepository(AssetRepository):
+    def _get_path(self, company_short_name: str, asset_type: AssetType, filename: str = "") -> Path:
+        return Path("companies") / company_short_name / asset_type.value / filename
+
+    def exists(self, company_short_name: str, asset_type: AssetType, filename: str) -> bool:
+        return self._get_path(company_short_name, asset_type, filename).is_file()
+
+    def read_text(self, company_short_name: str, asset_type: AssetType, filename: str) -> str:
+        path = self._get_path(company_short_name, asset_type, filename)
+        if not path.is_file():
+            raise FileNotFoundError(f"File not found: {path}")
+        return path.read_text(encoding="utf-8")
+
+    def list_files(self, company_short_name: str, asset_type: AssetType, extension: str = None) -> list[str]:
+        directory = self._get_path(company_short_name, asset_type)
+        if not directory.exists():
+            return []
+        files = [f.name for f in directory.iterdir() if f.is_file()]
+        if extension:
+            files = [f for f in files if f.endswith(extension)]
+        return files
+
+    def write_text(self, company_short_name: str, asset_type: AssetType, filename: str, content: str) -> None:
+        path = self._get_path(company_short_name, asset_type, filename)
+        # Ensure the directory exists (e.g. creating a new company structure)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(content, encoding="utf-8")
+
+    def delete(self, company_short_name: str, asset_type: AssetType, filename: str) -> None:
+        path = self._get_path(company_short_name, asset_type, filename)
+        if path.exists():
+            path.unlink()

iatoolkit/repositories/llm_query_repo.py

@@ -0,0 +1,105 @@
+# Copyright (c) 2024 Fernando Libedinsky
+# Product: IAToolkit
+#
+# IAToolkit is open source software.
+
+from iatoolkit.repositories.models import LLMQuery, Tool, 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 commit(self):
+        self.session.commit()
+
+    def rollback(self):
+        self.session.rollback()
+
+    def add_query(self, query: LLMQuery):
+        self.session.add(query)
+        self.session.commit()
+        return query
+
+
+    def get_company_tools(self, company: Company) -> list[Tool]:
+        return (
+            self.session.query(Tool)
+            .filter(
+                Tool.is_active.is_(True),
+                or_(
+                    Tool.company_id == company.id,
+                    Tool.system_function.is_(True)
+                )
+            )
+            # Ordenamos descendente: True (System) va primero, False (Company) va después
+            .order_by(Tool.system_function.desc())
+            .all()
+        )
+
+    def delete_system_tools(self):
+        self.session.query(Tool).filter_by(system_function=True).delete(synchronize_session=False)
+
+    def create_or_update_tool(self, new_tool: Tool):
+        tool = self.session.query(Tool).filter_by(company_id=new_tool.company_id,
+                                                  name=new_tool.name).first()
+        if tool:
+            tool.description = new_tool.description
+            tool.parameters = new_tool.parameters
+            tool.system_function = new_tool.system_function
+        else:
+            self.session.add(new_tool)
+            tool = new_tool
+
+        self.session.flush()
+        return tool
+
+    def delete_tool(self, tool: Tool):
+        self.session.query(Tool).filter_by(id=tool.id).delete(synchronize_session=False)
+
+    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.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.flush()
+        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.flush()
+        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_prompt_by_name(self, company: Company, prompt_name: str):
+        return self.session.query(Prompt).filter_by(company_id=company.id, name=prompt_name).first()
+
+    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()
+

iatoolkit/repositories/models.py

@@ -0,0 +1,279 @@
+# 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
+
+
+# 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('role', String, nullable=True, default='user'),
+                     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, autoincrement=True)
+    company_id = Column(Integer, ForeignKey('iat_companies.id', ondelete='CASCADE'), nullable=False)
+    key_name = Column(String, nullable=False)
+    key = Column(String, 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, autoincrement=True)
+    short_name = Column(String, nullable=False, unique=True, index=True)
+    name = Column(String, nullable=False)
+
+    parameters = Column(JSON, nullable=True)
+    created_at = Column(DateTime, default=datetime.now)
+
+    documents = relationship("Document",
+                             back_populates="company",
+                             cascade="all, delete-orphan",
+                             lazy='dynamic')
+    tools = relationship("Tool",
+                           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")
+
+    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, autoincrement=True)
+    email = Column(String, unique=True, nullable=False)
+    first_name = Column(String, nullable=False)
+    last_name = Column(String, 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, 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 Tool(Base):
+    """Represents a custom or system function that the LLM can call (tool)."""
+    __tablename__ = 'iat_tools'
+
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    company_id = Column(Integer,
+                        ForeignKey('iat_companies.id',ondelete='CASCADE'),
+                        nullable=True)
+    name = Column(String, 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='tools')
+
+    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, nullable=False, index=True)
+    meta = Column(JSON, nullable=True)
+    created_at = Column(DateTime, default=datetime.now)
+    content = Column(Text, nullable=False)
+    content_b64 = Column(Text, nullable=False)
+
+    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, autoincrement=True)
+    company_id = Column(Integer, ForeignKey('iat_companies.id',
+                            ondelete='CASCADE'), nullable=False)
+    user_identifier = Column(String, nullable=False)
+    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")
+
+    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, autoincrement=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)
+
+    # the size of this vector should be set depending on the embedding model used
+    # for OpenAI is 1536, and for huggingface is 384
+    embedding = Column(Vector(1536), nullable=False)
+
+    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 UserFeedback(Base):
+    """Stores feedback and ratings submitted by users for specific interactions."""
+    __tablename__ = 'iat_feedback'
+
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    company_id = Column(Integer, ForeignKey('iat_companies.id',
+                                            ondelete='CASCADE'), nullable=False)
+    user_identifier = Column(String, 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, autoincrement=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, autoincrement=True)
+    company_id = Column(Integer, ForeignKey('iat_companies.id',
+                                            ondelete='CASCADE'), nullable=True)
+    name = Column(String, nullable=False)
+    description = Column(String, nullable=False)
+    filename = Column(String, nullable=False)
+    active = Column(Boolean, default=True)
+    is_system_prompt = Column(Boolean, default=False)
+    order = Column(Integer, nullable=True, default=0)
+    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, autoincrement=True)
+
+    timestamp = Column(DateTime(timezone=True), server_default=func.now(), nullable=False, index=True)
+    company_short_name = Column(String, nullable=False, index=True)
+    user_identifier = Column(String, index=True)
+
+    # Cómo y el Resultado
+    auth_type = Column(String, nullable=False) # 'local', 'external_api', 'redeem_token', etc.
+    outcome = Column(String, nullable=False)   # 'success' o 'failure'
+    reason_code = Column(String)               # Causa de fallo, ej: 'INVALID_CREDENTIALS'
+
+    # Contexto de la Petición
+    source_ip = Column(String, nullable=False)
+    user_agent_hash = Column(String)           # Hash corto del User-Agent
+    request_path = Column(String, 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,171 @@
+# Copyright (c) 2024 Fernando Libedinsky
+# Product: IAToolkit
+#
+# IAToolkit is open source software.
+
+from iatoolkit.repositories.models import (User, Company, user_company,
+                                           ApiKey, UserFeedback, AccessLog)
+from injector import inject
+from iatoolkit.repositories.database_manager import DatabaseManager
+from sqlalchemy import select, func, and_
+
+
+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 get_user_role_in_company(self, company_id, user_id, ):
+        stmt = (
+            select(user_company.c.role)
+            .where(
+                user_company.c.user_id == user_id,
+                user_company.c.company_id == company_id,
+            )
+        )
+        result = self.session.execute(stmt).scalar_one_or_none()
+        return result
+
+    def get_admin_companies_by_user_identifier(self, user_identifier: str) -> list[Company]:
+        """
+        Return all the companies to which the user belongs (by email),
+        and where also he has admin role in the iat_user_company table.
+        """
+        return (
+            self.session.query(Company)
+            .join(user_company, Company.id == user_company.c.company_id)
+            .join(User, User.id == user_company.c.user_id)
+            .filter(User.email == user_identifier)
+            .filter(user_company.c.role == "admin")
+            .all()
+        )
+
+    def get_company_users_with_details(self, company_short_name: str) -> list[dict]:
+        # returns the list of users in the company with their role and last access date
+
+        # subquery for last access date
+        last_access_sq = (
+            self.session.query(
+                AccessLog.user_identifier,
+                func.max(AccessLog.timestamp).label("max_ts")
+            )
+            .filter(AccessLog.company_short_name == company_short_name)
+            .group_by(AccessLog.user_identifier)
+            .subquery()
+        )
+
+        # main query
+        stmt = (
+            self.session.query(
+                User,
+                user_company.c.role,
+                last_access_sq.c.max_ts
+            )
+            .join(user_company, User.id == user_company.c.user_id)
+            .join(Company, Company.id == user_company.c.company_id)
+            .outerjoin(last_access_sq, User.email == last_access_sq.c.user_identifier)
+            .filter(Company.short_name == company_short_name)
+        )
+
+        results = stmt.all()
+
+        return results
+
+    def create_company(self, new_company: Company):
+        company = self.session.query(Company).filter_by(short_name=new_company.short_name).first()
+        if company:
+            if company.parameters != new_company.parameters:
+                company.parameters = new_company.parameters
+        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.
+        """
+        api_key_entry = (self.session.query(ApiKey).filter
+                             (ApiKey.key == api_key_value, ApiKey.is_active == True).first())
+
+        return api_key_entry
+
+
+    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()
+
+
+
+