mfcli 0.2.1__py3-none-any.whl

mfcli/utils/datasheet_vectorizer.py

@@ -0,0 +1,283 @@
+import os
+import re
+from urllib.parse import urlparse, unquote
+
+from playwright.async_api import async_playwright, Browser
+from requests import RequestException
+from sqlmodel import select
+
+from mfcli.client.chroma_db import ChromaClient
+from mfcli.client.docling import DoclingChunker
+from mfcli.client.vector_db import DocumentVectorizer
+from mfcli.constants.file_types import PDFMimeTypes
+from mfcli.digikey.digikey import DigiKey
+from mfcli.models.bom import BOM
+from mfcli.models.datasheet import Datasheet
+from mfcli.pipeline.extractor import TextExtractor
+from mfcli.utils.directory_manager import app_dirs
+from mfcli.utils.http_requests import http_request
+from mfcli.utils.logger import get_logger
+from mfcli.utils.orm import Session
+from mfcli.utils.tools import get_mime_type_from_bytes
+
+logger = get_logger(__name__)
+
+
+class DatasheetVectorizer:
+    def __init__(self, chroma_db: ChromaClient):
+        self._extractor = TextExtractor()
+        self._vectorizer = DocumentVectorizer(chroma_db)
+        self._docling = DoclingChunker()
+
+    def _vectorize_text(self, text: str, file_name: str, purpose: str, additional_metadata: dict = None):
+        """
+        Shared method to vectorize text with metadata
+        :param text: Extracted text content
+        :param file_name: Name of the file
+        :param purpose: Purpose of the vectorization (e.g., 'datasheet', 'bom', 'errata')
+        :param additional_metadata: Optional additional metadata to include
+        """
+        metadata = {"file_name": file_name, "purpose": purpose}
+        if additional_metadata:
+            metadata.update(additional_metadata)
+        self._vectorizer.vectorize(text, metadata)
+        logger.debug(f"File vectorized: {file_name} (purpose: {purpose})")
+
+    @staticmethod
+    async def _fetch_with_playwright(browser: Browser, url: str):
+        context = await browser.new_context()
+        response = await context.request.get(url)
+        body = await response.body()
+        return body
+
+    @staticmethod
+    def _parse_ti_url(url: str) -> str:
+        """
+        Texas Instruments URLs may have goTo param which is the real URL of the PDF
+        :param url: TI URL
+        :return: URL from goTo param
+        """
+        url_query_params = urlparse(url).query
+        if not url_query_params:
+            return url
+        params = url_query_params.split('&')
+        for param in params:
+            name = param.split('=')[0]
+            value = param.split('=')[1]
+            if not name == 'gotoUrl':
+                continue
+            return unquote(value)
+        return url
+
+    @staticmethod
+    def _save_datasheet(name: str, content: bytes):
+        file_path = app_dirs.data_sheets_dir / name
+        with open(file_path, "wb") as f:
+            f.write(content)
+
+    async def _download(self, browser: Browser, url: str, purpose: str = "datasheet"):
+        logger.debug(f"Fetching datasheet: {url}")
+        try:
+            ti_url_regex = r"^https?://www.ti.com/.+$"
+            if re.match(ti_url_regex, url, re.I):
+                logger.debug(f"URL is a TI URL: {url}")
+                url = self._parse_ti_url(url)
+                logger.debug(f"Parsed TI URL: {url}")
+            url_path = urlparse(url).path
+        except ValueError as e:
+            logger.debug(f"Unable to parse datasheet URL: {url}")
+            logger.debug(e)
+            return
+        file_name = os.path.basename(url_path)
+        if not file_name.endswith(".pdf"):
+            file_name = f"{file_name}.pdf"
+        try:
+            content = http_request(method='GET', url=url).content
+            mime_type = get_mime_type_from_bytes(content, file_name)
+            if mime_type not in PDFMimeTypes:
+                logger.debug(f"Retrieved PDF is not PDF MIME type: {url}")
+                logger.debug(f"Retrying with playwright: {url}")
+                content = await self._fetch_with_playwright(browser, url)
+        except RequestException as e:
+            logger.debug(e)
+            logger.debug(f"HTTP error fetching PDF: {url}")
+            logger.debug(f"Retrying with playwright: {url}")
+            content = await self._fetch_with_playwright(browser, url)
+        except Exception as e:
+            logger.debug(f"Unhandled error fetching datasheet URL: {url}")
+            logger.debug(e)
+            return
+        mime_type = get_mime_type_from_bytes(content, file_name)
+        if mime_type not in PDFMimeTypes:
+            logger.debug(f"Could not fetch PDF even with playwright: {url}")
+            return
+        try:
+            self._save_datasheet(file_name, content)
+        except Exception as e:
+            logger.debug(e)
+            logger.debug(f"Error saving datasheet: {file_name}")
+
+    async def download(self, urls: list[str], purpose: str = "datasheet"):
+        if not urls:
+            logger.debug(f"No datasheets to vectorize, exiting")
+            return
+        logger.debug(f"Vectorizing {len(urls)} documents (purpose: {purpose})")
+        async with async_playwright() as p:
+            browser = await p.chromium.launch(headless=True)
+            try:
+                for url in urls:
+                    try:
+                        await self._download(browser, url, purpose)
+                    except Exception as e:
+                        logger.debug(e)
+                        logger.debug(f"Error processing document: {url}")
+            finally:
+                await browser.close()
+
+    def vectorize_local_file(self, file_path: str, purpose: str, additional_metadata: dict = None):
+        """
+        Vectorize a local file (e.g., generated by agents)
+        :param file_path: Path to the local file
+        :param purpose: Purpose of the vectorization (e.g., 'bom', 'errata', 'functional_blocks')
+        :param additional_metadata: Optional additional metadata to include
+        """
+        try:
+            logger.debug(f"Vectorizing local file: {file_path} (purpose: {purpose})")
+            file_name = os.path.basename(file_path)
+
+            # Check if file exists
+            if not os.path.exists(file_path):
+                logger.error(f"File does not exist: {file_path}")
+                return
+
+            # Extract text based on file type
+            with open(file_path, 'rb') as f:
+                content = f.read()
+
+            mime_type = get_mime_type_from_bytes(content, file_name)
+
+            if mime_type in PDFMimeTypes:
+                text = self._extractor.extract_pdf_bytes(content)
+            else:
+                # For non-PDF files, use the general extractor
+                text = self._extractor.extract_text_from_file_bytes(file_name, content)
+
+            logger.debug(f"Text extracted from local file: {file_path}")
+            self._vectorize_text(text, file_name, purpose, additional_metadata)
+
+        except Exception as e:
+            logger.error(f"Error vectorizing local file: {file_path}")
+            logger.exception(e)
+            raise
+
+    def vectorize_local_files(self, file_paths: list[str], purpose: str, additional_metadata: dict = None):
+        """
+        Vectorize multiple local files
+        :param file_paths: List of paths to local files
+        :param purpose: Purpose of the vectorization (e.g., 'bom', 'errata', 'functional_blocks')
+        :param additional_metadata: Optional additional metadata to include
+        """
+        if not file_paths:
+            logger.debug(f"No files to vectorize, exiting")
+            return
+
+        logger.debug(f"Vectorizing {len(file_paths)} local files (purpose: {purpose})")
+        for file_path in file_paths:
+            try:
+                self.vectorize_local_file(file_path, purpose, additional_metadata)
+            except Exception as e:
+                logger.exception(e)
+                logger.error(f"Error processing local file: {file_path}")
+        logger.debug(f"Finished vectorizing {len(file_paths)} local files")
+
+    def vectorize_text_content(self, text: str, file_name: str, purpose: str, additional_metadata: dict = None):
+        """
+        Vectorize text content directly (e.g., from agent output)
+        :param text: Text content to vectorize
+        :param file_name: Name to associate with this content
+        :param purpose: Purpose of the vectorization (e.g., 'bom', 'errata', 'functional_blocks')
+        :param additional_metadata: Optional additional metadata to include
+        """
+        try:
+            logger.debug(f"Vectorizing text content: {file_name} (purpose: {purpose})")
+            self._vectorize_text(text, file_name, purpose, additional_metadata)
+        except Exception as e:
+            logger.error(f"Error vectorizing text content: {file_name}")
+            logger.exception(e)
+            raise
+
+    def vectorize_file_buf(
+            self,
+            file_bytes: bytes,
+            file_name: str,
+            purpose: str,
+            additional_metadata: dict = None
+    ) -> None:
+        """
+        Vectorize a file from a buffer. This vectorizer uses DoclingChunker.
+        :param file_bytes: file bytes
+        :param file_name: file name
+        :param purpose: file purpose
+        :param additional_metadata: dict of metadata
+        :return: None
+        """
+        chunks = self._docling.chunk(file_name, file_bytes)
+        metadata = {"file_name": file_name, "purpose": purpose}
+        if additional_metadata:
+            metadata.update(additional_metadata)
+        self._vectorizer.vectorize_chunks(chunks, metadata)
+
+
+async def get_datasheets_for_bom_entries(db: Session, chroma_db: ChromaClient, entries: list[BOM]):
+    logger.info(f"Fetching datasheets for {len(entries)} BOM entries")
+    part_numbers = {entry.value for entry in entries}
+    logger.debug("Fetching existing datasheets for part numbers")
+
+    # Fetch existing datasheets
+    stmt = select(Datasheet).where(Datasheet.part_number.in_(part_numbers))
+    datasheets: list[Datasheet] = db.execute(stmt).scalars().all()
+    datasheet_map = {d.part_number: d.datasheet for d in datasheets}
+
+    logger.debug(f"Datasheet map: {datasheet_map}")
+    client = DigiKey()
+    new_datasheets: list[Datasheet] = []
+    datasheet_urls: list[str] = []
+    for entry in entries:
+        try:
+            logger.debug(f"Processing BOM entry: {entry.value}")
+
+            # Skip resistors, capacitors and inductors
+            ref = entry.reference
+            if ref.startswith('R') \
+            or ref.startswith('C') \
+            or ref.startswith('L') \
+            or ref.startswith('J') \
+            or ref.startswith('T') \
+            or ref.startswith('D'):
+                logger.debug(f"Skipping BOM entry {entry.value} with reference: {ref}")
+                continue
+
+            existing_datasheet = datasheet_map.get(entry.value)
+            if not existing_datasheet:
+                logger.debug(f"Datasheet does not exist for {entry.value}")
+            entry.datasheet = existing_datasheet or client.datasheet(entry.value)
+
+            # If datasheet is new create it in DB
+            if not existing_datasheet and entry.datasheet:
+                logger.debug(f"Adding new datasheet for {entry.value}: {entry.datasheet}")
+                new_datasheets.append(Datasheet(part_number=entry.value, datasheet=entry.datasheet))
+                datasheet_urls.append(entry.datasheet)
+        except Exception as e:
+            logger.error(f"Error adding datasheet for BOM entry: {entry.value}")
+            logger.exception(e)
+    if new_datasheets:
+        db.add_all(new_datasheets)
+        logger.debug(f"Adding new data sheets: {new_datasheets}")
+    if datasheet_urls:
+        logger.debug(f"About to download datasheets: {datasheet_urls}")
+        try:
+            await DatasheetVectorizer(chroma_db).download(datasheet_urls)
+        except Exception as e:
+            logger.error("Error vectorizing datasheets for BOM")
+            raise e
+    logger.debug("Finished adding datasheets")

mfcli/utils/directory_manager.py

@@ -0,0 +1,116 @@
+import os
+import sys
+from pathlib import Path
+from mfcli.utils.tools import get_git_root
+
+
+class DirectoryManager:
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            cls._instance._initialized = False
+        return cls._instance
+
+    def __init__(self):
+        if self._initialized:
+            return
+        # OS-specific base appdata location
+        if os.name == "nt":
+            app_data_base = Path(os.getenv("LOCALAPPDATA", os.getenv("APPDATA")))
+        elif sys.platform == "darwin":
+            app_data_base = Path.home() / "Library" / "Application Support"
+        else:
+            app_data_base = Path.home() / ".local" / "share"
+
+        self.home_dir: Path = Path(os.path.expanduser("~")) / "Multifactor"
+        self.env_file_path: Path = self.home_dir / ".env"
+
+        # User app directories
+        self.app_data_dir: Path = app_data_base / "Multifactor"
+        self.chroma_db_dir: Path = self.app_data_dir / "chromadb"
+
+        self.app_data_dir.mkdir(exist_ok=True, parents=True)
+        self.chroma_db_dir.mkdir(exist_ok=True, parents=True)
+        self.home_dir.mkdir(exist_ok=True, parents=True)
+
+        # Repo dirs
+        self.root_dir: Path | None = None
+        self.context_dir: Path | None = None
+        self.agent_instructions_dir: Path | None = None
+        self.data_sheets_dir: Path | None = None
+        self.fw_tasks_dir: Path | None = None
+        self.generated_files_dir: Path | None = None
+        self.cheat_sheets_dir: Path | None = None
+        self.reqs_dir: Path | None = None
+        self.pdf_parts_dir: Path | None = None
+        self.metadata_dir: Path | None = None
+        self.config_file_path: Path | None = None
+        self.file_docket_path: Path | None = None
+
+        self._initialized = True
+
+    def initialize(self, root: str):
+        # Accept file or directory
+        root_path = Path(root)
+        if root_path.is_file():
+            self.root_dir = root_path.parent
+        else:
+            self.root_dir = root_path
+
+        # Determine the base directory for project folders
+        # If in a git repo, use the git root; otherwise use the current directory
+        git_root = get_git_root(self.root_dir)
+        if git_root:
+            # Use git root for all project folders
+            base_dir = git_root
+        else:
+            # Not a git repo, use the current directory
+            base_dir = self.root_dir
+        
+        # Create "multifactor" parent folder at the base directory
+        multifactor_parent = base_dir / "multifactor"
+        
+        # Context folder - where users place files to be ingested
+        self.context_dir = multifactor_parent / "context"
+        
+        # Repo directories - all created within the "multifactor" folder
+        self.agent_instructions_dir = multifactor_parent / "agent_instructions"
+        self.data_sheets_dir = multifactor_parent / "data_sheets"
+        self.fw_tasks_dir = multifactor_parent / "fw_tasks"
+        self.generated_files_dir = multifactor_parent / "generated_files"
+        self.cheat_sheets_dir = multifactor_parent / "hw_cheat_sheets"
+        self.reqs_dir = multifactor_parent / "requirements"
+        self.pdf_parts_dir = multifactor_parent / "pdf_parts"
+        
+        # Metadata directory - also within the "multifactor" folder
+        self.metadata_dir = multifactor_parent
+        self.config_file_path = self.metadata_dir / "config.json"
+        self.file_docket_path = self.metadata_dir / "file_docket.json"
+
+        # Create all dirs
+        self._create_directory_structure()
+
+    def _create_directory_structure(self):
+        for directory in [
+            self.context_dir,
+            self.agent_instructions_dir,
+            self.data_sheets_dir,
+            self.fw_tasks_dir,
+            self.generated_files_dir,
+            self.cheat_sheets_dir,
+            self.reqs_dir,
+            self.pdf_parts_dir,
+            self.app_data_dir,
+            self.chroma_db_dir,
+            self.metadata_dir
+        ]:
+            directory.mkdir(parents=True, exist_ok=True)
+
+
+app_dirs = DirectoryManager()
+
+
+def init_directory_structure(root_dir: str):
+    app_dirs.initialize(root_dir)

mfcli/utils/file_upload.py

@@ -0,0 +1,298 @@
+"""Unified file upload abstraction for different LLM providers."""
+
+import os
+from abc import ABC, abstractmethod
+from enum import Enum
+from pathlib import Path
+from typing import Optional
+
+from google import genai
+from google.genai.types import File
+
+from mfcli.utils.config import get_config
+from mfcli.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class FileUploadProvider(str, Enum):
+    """Supported file upload providers."""
+    GEMINI = "gemini"
+    OPENAI = "openai"
+
+
+class BaseFileUploader(ABC):
+    """Base class for file upload implementations."""
+
+    @abstractmethod
+    def upload_file(self, file_path: str, display_name: Optional[str] = None) -> dict:
+        """
+        Upload a file to the provider's file storage.
+        
+        Args:
+            file_path: Path to the local file to upload
+            display_name: Optional display name for the file
+            
+        Returns:
+            Dictionary containing file metadata including URI/ID for accessing the file
+        """
+        pass
+
+    @abstractmethod
+    def delete_file(self, file_id: str) -> bool:
+        """
+        Delete a file from the provider's storage.
+        
+        Args:
+            file_id: The ID/URI of the file to delete
+            
+        Returns:
+            True if deletion was successful, False otherwise
+        """
+        pass
+
+    @abstractmethod
+    def get_file_info(self, file_id: str) -> dict:
+        """
+        Get information about an uploaded file.
+        
+        Args:
+            file_id: The ID/URI of the file
+            
+        Returns:
+            Dictionary containing file metadata
+        """
+        pass
+
+
+class GeminiFileUploader(BaseFileUploader):
+    """File uploader implementation for Google Gemini."""
+
+    def __init__(self):
+        """Initialize the Gemini file uploader."""
+        config = get_config()
+        self.client = genai.Client(api_key=config.google_api_key)
+        logger.info("Initialized Gemini file uploader")
+
+    @staticmethod
+    def _file_access_check(file_path: str):
+        file_path_obj = Path(file_path)
+
+        # Validate file exists and is readable
+        if not file_path_obj.exists():
+            raise ValueError(f"File does not exist: {file_path}")
+        if not os.access(file_path_obj, os.R_OK):
+            raise ValueError(f"File is not readable: {file_path}")
+
+    def upload(self, file_path: str) -> File:
+        """
+        Upload a file to Gemini Files API and return File object.
+
+        Args:
+            file_path: Path to the local file to upload
+
+        Returns:
+            Gemini types File object.
+
+        Raises:
+            ValueError: If file doesn't exist or is not readable
+            Exception: If upload fails
+        """
+        self._file_access_check(file_path)
+        return self.client.files.upload(file=file_path)
+
+    def upload_file(self, file_path: str, display_name: Optional[str] = None) -> dict:
+        """
+        Upload a file to Gemini Files API.
+        
+        Args:
+            file_path: Path to the local file to upload
+            display_name: Optional display name for the file
+            
+        Returns:
+            Dictionary with file metadata including 'uri', 'name', 'mime_type', 'size_bytes'
+            
+        Raises:
+            ValueError: If file doesn't exist or is not readable
+            Exception: If upload fails
+        """
+
+        self._file_access_check(file_path)
+
+        file_path_obj = Path(file_path)
+
+        # Use filename as display name if not provided
+        if display_name is None:
+            display_name = file_path_obj.name
+
+        try:
+            logger.info(f"Uploading file to Gemini: {file_path}")
+
+            # Upload the file
+            uploaded_file = self.client.files.upload(
+                file=str(file_path_obj),
+                config={'display_name': display_name}
+            )
+
+            # Extract metadata
+            result = {
+                'uri': uploaded_file.uri,
+                'name': uploaded_file.name,
+                'display_name': uploaded_file.display_name,
+                'mime_type': uploaded_file.mime_type,
+                'size_bytes': uploaded_file.size_bytes,
+                'state': uploaded_file.state.name,
+                'provider': FileUploadProvider.GEMINI.value
+            }
+
+            logger.info(f"Successfully uploaded file: {result['name']}")
+            return result
+
+        except Exception as e:
+            logger.error(f"Failed to upload file to Gemini: {e}")
+            raise Exception(f"Failed to upload file to Gemini: {str(e)}")
+
+    def delete_file(self, file_id: str) -> bool:
+        """
+        Delete a file from Gemini Files API.
+        
+        Args:
+            file_id: The name/ID of the file (e.g., 'files/abc123')
+            
+        Returns:
+            True if deletion was successful, False otherwise
+        """
+        try:
+            logger.info(f"Deleting file from Gemini: {file_id}")
+            self.client.files.delete(name=file_id)
+            logger.info(f"Successfully deleted file: {file_id}")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to delete file from Gemini: {e}")
+            return False
+
+    def get_file_info(self, file_id: str) -> dict:
+        """
+        Get information about an uploaded file.
+        
+        Args:
+            file_id: The name/ID of the file (e.g., 'files/abc123')
+            
+        Returns:
+            Dictionary containing file metadata
+        """
+        try:
+            logger.info(f"Getting file info from Gemini: {file_id}")
+            file_info = self.client.files.get(name=file_id)
+
+            result = {
+                'uri': file_info.uri,
+                'name': file_info.name,
+                'display_name': file_info.display_name,
+                'mime_type': file_info.mime_type,
+                'size_bytes': file_info.size_bytes,
+                'state': file_info.state.name,
+                'provider': FileUploadProvider.GEMINI.value
+            }
+
+            return result
+
+        except Exception as e:
+            logger.error(f"Failed to get file info from Gemini: {e}")
+            raise Exception(f"Failed to get file info: {str(e)}")
+
+
+class OpenAIFileUploader(BaseFileUploader):
+    """File uploader implementation for OpenAI (placeholder for future implementation)."""
+
+    def __init__(self):
+        """Initialize the OpenAI file uploader."""
+        config = get_config()
+        # This will be implemented when OpenAI support is added
+        logger.info("OpenAI file uploader - not yet implemented")
+        raise NotImplementedError("OpenAI file upload support coming soon")
+
+    def upload_file(self, file_path: str, display_name: Optional[str] = None) -> dict:
+        """Upload a file to OpenAI."""
+        raise NotImplementedError("OpenAI file upload not yet implemented")
+
+    def delete_file(self, file_id: str) -> bool:
+        """Delete a file from OpenAI."""
+        raise NotImplementedError("OpenAI file deletion not yet implemented")
+
+    def get_file_info(self, file_id: str) -> dict:
+        """Get file info from OpenAI."""
+        raise NotImplementedError("OpenAI file info not yet implemented")
+
+
+class FileUploadManager:
+    """Manager class to handle file uploads across different providers."""
+
+    def __init__(self, provider: FileUploadProvider = FileUploadProvider.GEMINI):
+        """
+        Initialize the file upload manager.
+        
+        Args:
+            provider: The file upload provider to use (default: GEMINI)
+        """
+        self.provider = provider
+        self.uploader = self._get_uploader(provider)
+
+    def _get_uploader(self, provider: FileUploadProvider) -> BaseFileUploader:
+        """Get the appropriate uploader for the specified provider."""
+        if provider == FileUploadProvider.GEMINI:
+            return GeminiFileUploader()
+        elif provider == FileUploadProvider.OPENAI:
+            return OpenAIFileUploader()
+        else:
+            raise ValueError(f"Unsupported file upload provider: {provider}")
+
+    def upload_file(self, file_path: str, display_name: Optional[str] = None) -> dict:
+        """
+        Upload a file using the configured provider.
+        
+        Args:
+            file_path: Path to the local file to upload
+            display_name: Optional display name for the file
+            
+        Returns:
+            Dictionary with file metadata
+        """
+        return self.uploader.upload_file(file_path, display_name)
+
+    def delete_file(self, file_id: str) -> bool:
+        """
+        Delete a file using the configured provider.
+        
+        Args:
+            file_id: The ID/URI of the file to delete
+            
+        Returns:
+            True if deletion was successful, False otherwise
+        """
+        return self.uploader.delete_file(file_id)
+
+    def get_file_info(self, file_id: str) -> dict:
+        """
+        Get file information using the configured provider.
+        
+        Args:
+            file_id: The ID/URI of the file
+            
+        Returns:
+            Dictionary containing file metadata
+        """
+        return self.uploader.get_file_info(file_id)
+
+
+def get_file_upload_manager(provider: FileUploadProvider = FileUploadProvider.GEMINI) -> FileUploadManager:
+    """
+    Factory function to get a file upload manager instance.
+    
+    Args:
+        provider: The file upload provider to use (default: GEMINI)
+        
+    Returns:
+        FileUploadManager instance
+    """
+    return FileUploadManager(provider)