mfcli 0.2.0__py3-none-any.whl

mfcli/pipeline/preprocessors/user_guide.py

@@ -0,0 +1,127 @@
+import asyncio
+from pathlib import Path
+from typing import List
+
+from google.genai.types import File as GeminiFile
+from pydantic import BaseModel, Field
+
+from mfcli.agents.tools.general import format_instructions
+from mfcli.models.file import File
+from mfcli.models.pdf_parts import PDFPart
+from mfcli.pipeline.run_context import PipelineRunContext
+from mfcli.utils.directory_manager import app_dirs
+from mfcli.utils.pdf_splitter import PDFSplitter
+
+
+class TOCSection(BaseModel):
+    title: str = Field(..., description="Section title")
+    section_no: int = Field(..., description="Section number")
+    start_page: int = Field(..., description="Starting page")
+    end_page: int = Field(..., description="End page")
+
+
+class TOC(BaseModel):
+    sections: List[TOCSection] = Field(..., description="Table of Contents sections")
+
+
+user_guide_preprocessor_instructions = format_instructions(
+    """
+    You will receive the start of a PDF for hardware engineering user guide. 
+    Your task is to extract all the Table of Contents sections from the PDF.
+    You will respond with the section title (no numbers in the title).
+    You will respond separately with the section number (section_no).
+    You will also respond with the start page (start_page) and (end_page) of this section. 
+    You MUST respond with all relevant top-level sections in the PDF. 
+    
+    Here are examples of relevant top-level sections:
+        
+        1. Architecture
+        2. PMCU
+        3. CPU
+        
+    Here are examples of sections which are NOT relevant (content) sections (do not include these):
+    
+        Read This First
+        About This Manual
+        Glossary
+        Related Documentation
+        Support Resources 
+        
+    Here are examples of sections which are NOT top-level (do not include these):
+    
+        1.1 Architecture Overview
+        1.2 Bus Organization
+        1.3 Platform Memory Map
+        
+    ONLY include content sections and top-level sections. 
+    """
+)
+
+
+class UserGuidePreprocessor:
+    def __init__(
+            self,
+            context: PipelineRunContext,
+            file: File,
+            pdf_head: GeminiFile,
+            content: bytes,
+            splitter: PDFSplitter
+    ):
+        self._context = context
+        self._file = file
+        self._pdf_head = pdf_head
+        self._content = content
+        self._splitter = splitter
+
+    async def _generate_toc(self) -> TOC:
+        return await self._context.gemini.generate(
+            prompt="Generate the Table of Content sections for this PDF",
+            instructions=user_guide_preprocessor_instructions,
+            response_model=TOC,
+            files=[self._pdf_head]
+        )
+
+    async def _create_pdf_part(self, section: TOCSection, pdf_part_path: Path) -> PDFPart:
+        pdf_part_gemini_file = await self._context.gemini.upload(pdf_part_path)
+        self._context.gemini_file_cache[pdf_part_gemini_file.name] = pdf_part_gemini_file
+        return PDFPart(
+            path=str(pdf_part_path),
+            file_id=self._file.id,
+            gemini_file_id=pdf_part_gemini_file.name,
+            start_page=section.start_page,
+            end_page=section.end_page,
+            title=section.title,
+            section_no=section.section_no
+        )
+
+    async def preprocess(self) -> List[PDFPart]:
+        toc = await self._generate_toc()
+        upload_tasks = []
+        for section in toc.sections:
+            pdf_part_path = self._splitter.extract_range(
+                start_page=section.start_page,
+                end_page=section.end_page,
+                output_folder=app_dirs.pdf_parts_dir
+            )
+            upload_tasks.append(self._create_pdf_part(section, pdf_part_path))
+        pdf_parts: List[PDFPart] = await asyncio.gather(*upload_tasks)
+        return pdf_parts
+
+
+async def preprocess_user_guide(
+        context: PipelineRunContext,
+        file: File,
+        pdf_head: GeminiFile,
+        content: bytes,
+        splitter: PDFSplitter
+) -> None:
+    preprocessor = UserGuidePreprocessor(
+        context=context,
+        file=file,
+        pdf_head=pdf_head,
+        content=content,
+        splitter=splitter
+    )
+    pdf_parts = await preprocessor.preprocess()
+    context.db.add_all(pdf_parts)
+    context.db.commit()

mfcli/pipeline/run_context.py

@@ -0,0 +1,32 @@
+from typing import Dict
+
+from google.genai.types import File as GeminiFile
+from mfcli.utils.query_service import QueryService
+
+from mfcli.client.gemini import Gemini
+from mfcli.models.file_docket import FileDocket
+from mfcli.models.pipeline_run import PipelineRun
+from mfcli.models.project_metadata import ProjectConfig
+from mfcli.utils.datasheet_vectorizer import DatasheetVectorizer
+from mfcli.utils.orm import Session
+
+
+class PipelineRunContext:
+    def __init__(
+            self,
+            db: Session,
+            pipeline_run: PipelineRun,
+            gemini: Gemini,
+            gemini_file_cache: Dict[str, GeminiFile],
+            docket: FileDocket,
+            config: ProjectConfig,
+            vectorizer: DatasheetVectorizer
+    ):
+        self.db = db
+        self.run = pipeline_run
+        self.gemini = gemini
+        self.gemini_file_cache = gemini_file_cache
+        self.docket = docket
+        self.config = config
+        self.vectorizer = vectorizer
+        self.query_service = QueryService(self.db)

mfcli/pipeline/schema_mapper.py

@@ -0,0 +1,89 @@
+from typing import Type, Optional
+
+from pydantic import BaseModel, Field
+from sqlmodel import SQLModel
+
+from mfcli.agents.tools.general import format_instructions
+from mfcli.client.gemini import Gemini
+from mfcli.constants.file_types import FileSubtypes
+from mfcli.models.bom import BOM, BOMSchema
+from mfcli.models.netlist import Netlist
+from mfcli.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+SubtypeSchemas: dict[FileSubtypes, Type[BaseModel]] = {
+    FileSubtypes.BOM: BOMSchema
+}
+SubtypeModels: dict[FileSubtypes, Type[SQLModel]] = {
+    FileSubtypes.BOM: BOM,
+    FileSubtypes.KICAD_LEGACY_NET: Netlist,
+    FileSubtypes.PADS_PCB_ASCII: Netlist,
+    FileSubtypes.KICAD_SPICE: Netlist,
+    FileSubtypes.PROTEL_ALTIUM: Netlist
+}
+
+
+class SchemaMapping(BaseModel):
+    input_field: Optional[str] = Field(None, description="A field found in the sample file")
+    mapped_field: str = Field(..., description="The field found in the backend to be mapped to")
+
+
+class SchemaMappings(BaseModel):
+    fields: list[SchemaMapping] = Field(default_factory=list, description="List of schema mappings")
+
+
+schema_mapper_instructions = format_instructions(
+    """
+    You are responsible for mapping fields found in a file to a backend schema.
+    You will be given the schema for a file subtype, like Bill of Materials (BOM). 
+    You will map whatever fields you see in the file, to fields in the backend. 
+    For example, a column "Designator" found in a BOM file would map to "reference" in the backend.
+    
+    You must respond **only** with valid JSON that exactly matches the `SchemaMappings` model:
+    
+    - The top-level object must have a key `"fields"` containing a list.
+    - Each item in the list must be an object with:
+        - `"input_field"` (optional string) — the field name from the file header.
+        - `"mapped_field"` (required string) — the corresponding backend field name.
+    
+    Do **not** include any markdown, text, or code fences. Respond **only** with JSON.
+    
+    Example of valid response:
+    
+    {
+        "fields": [
+            {"input_field": "RefDes", "mapped_field": "reference"},
+            {"input_field": "Value", "mapped_field": "value"},
+            {"input_field": "Qty", "mapped_field": "quantity"},
+            {"input_field": "Description", "mapped_field": "description"}
+        ]
+    }
+    """
+)
+
+
+async def map_schema(gemini: Gemini, subtype: int, text: str) -> SchemaMappings | None:
+    if not SubtypeSchemas.get(subtype):
+        logger.debug(f"No subtype mapping required for subtype: {subtype}")
+        return
+    schema = str(SubtypeSchemas[subtype].model_json_schema())
+    prompt = format_instructions(
+        f"""
+        {schema_mapper_instructions}
+        
+        Here is the backend schema for this filetype:
+        
+        {schema}
+        
+        Here is the file contents:
+        
+        {text}
+        
+        """
+    )
+    return await gemini.generate(
+        prompt=prompt,
+        instructions=schema_mapper_instructions,
+        response_model=SchemaMappings
+    )

mfcli/pipeline/sub_classifier.py

@@ -0,0 +1,115 @@
+from typing import Literal
+
+from google.genai.types import File as GeminiFile
+from pydantic import BaseModel
+
+from mfcli.agents.tools.general import format_instructions
+from mfcli.client.gemini import Gemini
+from mfcli.constants.file_types import (
+    FileTypes,
+    FileSubtypes,
+    FILE_SUBTYPE_UNKNOWN,
+    PDFSubtypeDescriptions,
+    OtherFileTypeDescriptions,
+    PDFFileSubtypeNames,
+    OtherFileSubtypeNames
+)
+from mfcli.models.file import File
+from mfcli.pipeline.parsers.netlist.kicad_legacy_net import is_kicad_legacy_netlist
+from mfcli.pipeline.parsers.netlist.protel_detector import is_protel_netlist
+from mfcli.utils.files import is_text_mime_type
+from mfcli.utils.logger import get_logger
+
+sub_classifier_instructions = format_instructions(
+    """
+    You are the sub-classifier agent for an engineering document processing pipeline. 
+    You will receive the first 50 lines of text from a file.
+    You will examine the content of the file, and determine the sub-type of this file. 
+    You will be given the sub-types and sub-type descriptions. 
+    If you are not able to determine the sub-type you will respond with "UNKNOWN". 
+    
+    Here are the valid sub-types and a description of each:
+
+    {}
+    
+    """
+)
+
+logger = get_logger(__name__)
+
+
+class PDFSubtypeClassifierResponse(BaseModel):
+    type: PDFFileSubtypeNames
+
+
+class OtherFileSubtypeClassifierResponse(BaseModel):
+    type: OtherFileSubtypeNames
+
+
+FileClass = Literal['pdf', 'other']
+
+
+class FileSubtypeAnalyzer:
+    def __init__(self, gemini: Gemini):
+        self._gemini = gemini
+
+    async def _get_subtype_from_gemini(
+            self,
+            prompt: str,
+            instructions: str,
+            gemini_file: GeminiFile | None = None,
+            file_class: FileClass = 'other'
+    ) -> str:
+        model = OtherFileSubtypeClassifierResponse if file_class == 'other' else PDFSubtypeClassifierResponse
+        files = [gemini_file] if gemini_file else None
+        response = await self._gemini.generate(
+            prompt=prompt,
+            instructions=instructions,
+            response_model=model,
+            files=files
+        )
+        return response.type
+
+    async def _get_subtype(
+            self,
+            prompt: str,
+            file: File,
+            gemini_file: GeminiFile | None = None,
+            file_class: FileClass = 'other'
+    ) -> None:
+        logger.debug(f"Fetching subtype for file: {file.name}")
+        relevant_subtype_descriptions = PDFSubtypeDescriptions if file.type == FileTypes.PDF else OtherFileTypeDescriptions
+        logger.debug(f"Relevant subtypes: {relevant_subtype_descriptions.keys()}")
+        instructions = sub_classifier_instructions.format(relevant_subtype_descriptions)
+        subtype = await self._get_subtype_from_gemini(prompt, instructions, gemini_file, file_class)
+        logger.debug(f"Subtype discovered: {subtype}")
+        if subtype == FILE_SUBTYPE_UNKNOWN:
+            raise RuntimeError(f"Could not determine the file subtype for file: {file.name}")
+        if not subtype in relevant_subtype_descriptions:
+            raise RuntimeError(f"LLM responded with invalid subtype: {subtype}")
+        file.sub_type = FileSubtypes.get(subtype)
+
+    async def analyze_pdf(
+            self,
+            file: File,
+            gemini_file: GeminiFile
+    ) -> None:
+        prompt = "Determine this PDF file subtype"
+        await self._get_subtype(prompt, file, gemini_file, 'pdf')
+
+    async def analyze_file(
+            self,
+            file: File,
+            text: str
+    ) -> None:
+        # Handle text MIME types
+        if text and is_text_mime_type(file.mime_type):
+            if file.type == FileTypes.NET:
+                if is_kicad_legacy_netlist(text):
+                    file.sub_type = FileSubtypes.KICAD_LEGACY_NET.value
+                elif is_protel_netlist(text):
+                    file.sub_type = FileSubtypes.PROTEL_ALTIUM.value
+
+        # If subtype cannot be parsed, use LLM to determine subtype
+        if not file.sub_type:
+            await self._get_subtype(text[0:500], file)

mfcli/utils/config.py

@@ -0,0 +1,33 @@
+from functools import lru_cache
+
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from mfcli.constants.openai import OPENAI_DEFAULT_EMBEDDING_MODEL, OPENAI_DEFAULT_EMBEDDING_DIMENSIONS
+from mfcli.utils.directory_manager import app_dirs
+
+
+class Settings(BaseSettings):
+    digikey_client_id: str
+    digikey_client_secret: str
+    openai_api_key: str
+    google_api_key: str
+    log_level: str = Field(default="INFO")
+    use_docling: bool = Field(default=True)
+    chunk_tokens: int = Field(default=2000)
+    chunk_size: int = Field(default=500)
+    chunk_overlap: int = Field(default=50)
+    embedding_model: str = Field(default=OPENAI_DEFAULT_EMBEDDING_MODEL)
+    embedding_dimensions: int = Field(default=OPENAI_DEFAULT_EMBEDDING_DIMENSIONS)
+
+    model_config = SettingsConfigDict(
+        env_file=str(app_dirs.env_file_path),
+        extra="allow",
+        env_file_encoding="utf-8",
+        case_sensitive=False
+    )
+
+
+@lru_cache
+def get_config() -> Settings:
+    return Settings()