mfcli 0.2.1__py3-none-any.whl

mfcli/pipeline/pipeline.py

@@ -0,0 +1,470 @@
+import json
+import os.path
+from pathlib import Path
+from typing import Dict, List
+
+from google.genai.types import File as GeminiFile
+
+from mfcli.models.pdf_parts import PDFPart
+from mfcli.models.project_metadata import ProjectConfig
+from mfcli.pipeline.preprocessors.user_guide import preprocess_user_guide
+from mfcli.pipeline.run_context import PipelineRunContext
+from mfcli.utils.datasheet_vectorizer import DatasheetVectorizer
+
+from mfcli.agents.tools.general import format_error_for_llm
+from mfcli.client.chroma_db import ChromaClient
+from mfcli.client.gemini import Gemini
+from mfcli.constants.directory_names import MF_PROJECT_CONFIG_DIR_NAME
+from mfcli.constants.file_types import (
+    SchemalessFileSubtypes,
+    FileTypes,
+    FileSubtypes,
+    PDFNoVectorizeFileSubtypes,
+    SummaryCheatSheetSubtypes
+)
+from mfcli.crud.file import create_file
+from mfcli.crud.pipeline_run import create_pipeline_run
+from mfcli.crud.project import get_project_by_name, read_project_config_file
+from mfcli.models.file import File
+from mfcli.models.file_docket import FileDocket, FileDocketEntry
+from mfcli.models.pipeline_run import PipelineRun
+from mfcli.models.project import Project
+from mfcli.pipeline.analysis.bom_netlist_mapper import map_netlist_to_bom_entries
+from mfcli.pipeline.analysis.generators.generator import Generator
+from mfcli.pipeline.classifier import get_file_metadata, validate_file
+from mfcli.pipeline.data_enricher import enrich_data_for_model
+from mfcli.pipeline.extractor import extract_document_text
+from mfcli.pipeline.parser import parse_schema
+from mfcli.pipeline.schema_mapper import map_schema
+from mfcli.pipeline.sub_classifier import FileSubtypeAnalyzer
+from mfcli.utils.directory_manager import app_dirs
+from mfcli.utils.logger import get_logger
+from mfcli.utils.orm import Session
+from mfcli.utils.pdf_splitter import PDFSplitter
+
+logger = get_logger(__name__)
+
+
+# TODO: IMPROVE get_file_subtype SO IT DOESN'T USE LLM
+
+
+class PipelineRunner:
+    def __init__(self, db: Session, project: Project, project_config: ProjectConfig):
+        self._db = db
+        self._project = project
+        # Use context folder for file ingestion by default
+        self.folder_path = str(app_dirs.context_dir)
+        self.total_files = 0
+        self.successfully_processed = 0
+        self.failed_files = 0
+        self.skipped_files = 0
+        self.errors = []
+        self.pipeline_run: PipelineRun | None = None
+        self._gemini = Gemini()
+        self._gemini_file_cache: Dict[str, GeminiFile] = {}
+        self._chroma_db = ChromaClient(project.index_id)
+        self._docket = FileDocket()
+        self._vectorizer = DatasheetVectorizer(self._chroma_db)
+        self._subtype_analyzer = FileSubtypeAnalyzer(self._gemini)
+        self._config = project_config
+        self._context = PipelineRunContext(
+            db=self._db,
+            pipeline_run=self.pipeline_run,
+            gemini=self._gemini,
+            gemini_file_cache=self._gemini_file_cache,
+            docket=self._docket,
+            config=self._config,
+            vectorizer=self._vectorizer
+        )
+        # Track which file types were actually processed (not skipped) in this run
+        self._processed_file_types: set[str] = set()
+        # Load existing file docket if it exists
+        self._load_existing_docket()
+    
+    def _load_existing_docket(self):
+        """Load existing file docket from JSON file"""
+        if app_dirs.file_docket_path and app_dirs.file_docket_path.exists():
+            logger.info(f"Loading existing file docket from: {app_dirs.file_docket_path}")
+            self._docket.load_from_json(app_dirs.file_docket_path)
+        else:
+            logger.info("No existing file docket found, starting fresh")
+
+    def _add_to_file_docket(self, file: File):
+        if file.is_datasheet:
+            vectorize = self._config.vectorize_datasheets
+        else:
+            vectorize = self._config.vectorize_hw_files
+        entry = FileDocketEntry(
+            name=file.name,
+            path=file.path,
+            vectorize=vectorize,
+            sub_type=FileSubtypes(file.sub_type).name,
+            md5=file.md5,
+            is_datasheet=bool(file.is_datasheet)
+        )
+        self._docket.add(entry)
+
+    def _save_file_docket(self):
+        json_data = json.dumps(self._docket.get_entries(), indent=2)
+        with open(app_dirs.file_docket_path, "w") as f:
+            f.write(json_data)
+
+    def _check_context_folder_has_files(self) -> bool:
+        """
+        Check if the context folder contains any files.
+        Returns True if files exist, False otherwise.
+        """
+        context_path = Path(self.folder_path)
+        
+        if not context_path.exists():
+            logger.warning(f"Context folder does not exist: {context_path}")
+            return False
+        
+        # Check for any files in the context folder (recursively)
+        ignore_dirs = [MF_PROJECT_CONFIG_DIR_NAME]
+        for dir_path, dir_names, file_names in os.walk(context_path):
+            dir_names[:] = [d for d in dir_names if d not in ignore_dirs]
+            if file_names:
+                return True
+        
+        return False
+
+    def _display_empty_context_message(self):
+        """
+        Display a helpful message when the context folder is empty.
+        """
+        context_path = Path(self.folder_path)
+        
+        print(f"\n{'='*70}")
+        print(f"CONTEXT FOLDER IS EMPTY")
+        print(f"{'='*70}")
+        print(f"\nThe context folder contains no files to process:")
+        print(f"  {context_path}")
+        print(f"\nTo run the pipeline, please add critical files to this folder, such as:")
+        print(f"  • Bill of Materials (BOM) files")
+        print(f"  • Schematics (PDF or other supported formats)")
+        print(f"  • MCU/IC user manuals and datasheets")
+        print(f"  • Netlist files")
+        print(f"  • Reference designs")
+        print(f"  • Application notes")
+        print(f"  • Any other hardware design documentation")
+        print(f"\nOnce you've added your files, run 'mfcli run' again.")
+        print(f"{'='*70}\n")
+        
+        logger.info("Pipeline execution cancelled: context folder is empty")
+
+    async def _gemini_files_upload(self, files: List[File | PDFPart]) -> List[GeminiFile]:
+        gemini_files = []
+        for file in files:
+            logger.debug(f"Checking for {type(file)} Gemini file")
+            if file.gemini_file_id in self._gemini_file_cache:
+                gemini_files.append(self._gemini_file_cache[file.gemini_file_id])
+                continue
+            logger.debug(f"Uploading {type(file)} to Gemini API")
+            gemini_file = await self._gemini.upload(file.path)
+            self._gemini_file_cache[gemini_file.name] = gemini_file
+            file.gemini_file_id = gemini_file.name
+            gemini_files.append(gemini_file)
+        return gemini_files
+
+    async def _preprocess_pdf(self, file: File, file_path: str, content: bytes):
+        logger.debug(f"Uploading file to Gemini: {file_path}")
+        splitter = PDFSplitter(file.name, content)
+        logger.debug("Splitting PDF head")
+        pdf_head_path = splitter.split_pdf_head()
+        gemini_pdf_head_file = await self._gemini.upload(pdf_head_path)
+        logger.debug(f"Analyzing PDF subtype: {file_path}")
+        await self._subtype_analyzer.analyze_pdf(file, gemini_pdf_head_file)
+
+        # Pre-process user guide files which are too big to upload to Gemini normally
+        # Extract the table of contents and split PDF into relevant content sections
+        # These sections will be used to generate summaries in analysis phase of pipeline
+        if file.sub_type in SummaryCheatSheetSubtypes:
+            await preprocess_user_guide(
+                context=self._context,
+                file=file,
+                pdf_head=gemini_pdf_head_file,
+                content=content,
+                splitter=splitter
+            )
+        else:
+            gemini_files = await self._gemini_files_upload([file])
+            logger.debug(f"Gemini files: {gemini_files}")
+
+        if self._config.vectorize_hw_files:
+            if file.sub_type in PDFNoVectorizeFileSubtypes:
+                logger.debug(f"PDF subtype does not require vectorization, skipping")
+            else:
+                logger.info(f"Chunking and vectorizing PDF: {file_path}")
+
+                # Use Docling to chunk PDF file and then vectorize it
+                self._vectorizer.vectorize_file_buf(
+                    file_name=file.name,
+                    file_bytes=content,
+                    purpose=FileSubtypes(file.sub_type).name
+                )
+        else:
+            logger.debug(f"vectorize_hw_files is set to False, skipping")
+
+    async def _preprocess_other_file_types(self, file: File, file_path: str, content: bytes):
+        logger.debug(f"File is not a PDF: {file_path}")
+        text_content = extract_document_text(file, content)
+        logger.debug(f"Analyzing file subtype: {file_path}")
+
+        # Analyze subtype
+        await self._subtype_analyzer.analyze_file(file, text_content)
+
+        if self._config.vectorize_hw_files:
+            logger.info(f"Vectorizing file: {file_path}")
+
+            # Vectorize
+            self._vectorizer.vectorize_text_content(
+                text=text_content,
+                file_name=file.name,
+                purpose=FileSubtypes(file.sub_type).name,
+                additional_metadata={"is_datasheet": file.is_datasheet}
+            )
+
+        # Ignore file subtypes that have no schema to parse like schematic files
+        if file.sub_type in SchemalessFileSubtypes:
+            logger.debug(f"File subtype is in ignore list, not parsing")
+        else:
+            logger.debug(f"File subtype is not in ignore list, parsing")
+
+            # Map schema
+            logger.debug(f"Mapping schema: {file_path}")
+            schema_mapping = await map_schema(self._gemini, file.sub_type, text_content)
+
+            # Parse schema from file
+            instances = parse_schema(self._db, file, schema_mapping)
+
+            # Enrich data
+            await enrich_data_for_model(self._db, self._chroma_db, file.sub_type, instances)
+
+    async def _preprocess_file(self, file_path: str, is_datasheet: bool) -> File:
+        logger.debug(f"Initializing Gemini client for pipeline: {self.pipeline_run.id}")
+
+        # Retrieve metadata
+        logger.debug(f"Retrieving file metadata: {file_path}")
+        metadata, content = get_file_metadata(file_path, is_datasheet)
+        logger.debug(f"Metadata retrieved: {file_path}")
+
+        # Validate file
+        validate_file(metadata)
+        logger.debug(f"File type validated: {file_path}")
+
+        # Create file
+        file = create_file(self._db, self.pipeline_run.id, metadata)
+
+        # PDF file pre-processing
+        if file.type == FileTypes.PDF:
+            await self._preprocess_pdf(file, file_path, content)
+
+        # Other file types pre-processing
+        else:
+            await self._preprocess_other_file_types(file, file_path, content)
+
+        self._db.commit()
+        logger.info(f"Pre-processing complete: {file_path}")
+        return file
+
+    async def _run_analysis(self):
+        logger.debug(f"Starting analysis for pipeline: {self.pipeline_run.id}")
+        
+        # Only run netlist-to-BOM mapping if BOM or netlist files were processed
+        netlist_types = {"PROTEL_ALTIUM", "KICAD_LEGACY_NET", "KICAD_SPICE", "PADS", "EDIF"}
+        should_map_netlist = (
+            "BOM" in self._processed_file_types or 
+            any(nt in self._processed_file_types for nt in netlist_types)
+        )
+        
+        if should_map_netlist:
+            logger.info("Running netlist-to-BOM mapping (BOM or netlist files were processed)")
+            map_netlist_to_bom_entries(self._db, self.pipeline_run.id)
+        else:
+            logger.info("Skipping netlist-to-BOM mapping (no BOM or netlist files processed)")
+        
+        logger.debug(f"Finished analysis for pipeline: {self.pipeline_run.id}")
+        
+        # Pass processed file types to generator for conditional generation
+        generator = Generator(self._context, self._processed_file_types)
+        await generator.generate_cheat_sheets()
+
+    def _check_file_in_docket(self, file_path: str, file_md5: str) -> tuple[bool, bool]:
+        """
+        Check if file exists in docket and if MD5 matches.
+        Returns: (should_process, should_remove_old)
+        """
+        # Check if file with same path exists in docket
+        existing_entry = self._docket.get_by_path(file_path)
+        
+        if not existing_entry:
+            # File not in docket, process it
+            return True, False
+        
+        # File exists in docket, check MD5
+        if existing_entry.md5 == file_md5:
+            # MD5 matches, skip processing
+            logger.info(f"File already processed with matching MD5, skipping: {file_path}")
+            return False, False
+        
+        # MD5 doesn't match, prompt user
+        logger.warning(f"File exists in docket but MD5 has changed: {file_path}")
+        logger.warning(f"  Old MD5: {existing_entry.md5}")
+        logger.warning(f"  New MD5: {file_md5}")
+        
+        print(f"\n{'='*70}")
+        print(f"File has been modified: {Path(file_path).name}")
+        print(f"Path: {file_path}")
+        print(f"Old MD5: {existing_entry.md5}")
+        print(f"New MD5: {file_md5}")
+        print(f"{'='*70}")
+        response = input("Do you want to delete the old file data and process the new version? (yes/no): ").strip().lower()
+        
+        if response in ['yes', 'y']:
+            logger.info(f"User confirmed deletion and reprocessing of: {file_path}")
+            return True, True
+        else:
+            logger.info(f"User declined reprocessing, skipping: {file_path}")
+            return False, False
+
+    def _remove_file_from_kb(self, entry: FileDocketEntry):
+        """Remove file from knowledge base (ChromaDB)"""
+        try:
+            collection = self._chroma_db._collection
+            results = collection.get()
+            
+            if not results or not results.get('metadatas'):
+                logger.warning(f"No data found in knowledge base to remove for: {entry.name}")
+                return
+            
+            # Find matching chunks for this file
+            matching_ids = []
+            for idx, metadata in enumerate(results['metadatas']):
+                if metadata and metadata.get('file_name') == entry.name:
+                    matching_ids.append(results['ids'][idx])
+            
+            if matching_ids:
+                collection.delete(ids=matching_ids)
+                logger.info(f"Removed {len(matching_ids)} chunks from knowledge base for: {entry.name}")
+            else:
+                logger.info(f"No chunks found in knowledge base for: {entry.name}")
+                
+            # Remove from docket
+            self._docket.remove(entry)
+            logger.info(f"Removed file from docket: {entry.name}")
+            
+        except Exception as e:
+            logger.error(f"Error removing file from knowledge base: {entry.name}")
+            logger.exception(e)
+
+    async def _preprocess_folder(self, folder_path: str | Path, is_datasheet: bool):
+        ignore_dirs = [MF_PROJECT_CONFIG_DIR_NAME]
+        for dir_path, dir_names, file_names in os.walk(folder_path):
+            dir_names[:] = [d for d in dir_names if d not in ignore_dirs]
+            for file_name in file_names:
+                self.total_files += 1
+                file_path = os.path.join(dir_path, file_name)
+                
+                try:
+                    # Get file metadata to check MD5
+                    logger.debug(f"Checking file: {file_path}")
+                    metadata, _ = get_file_metadata(file_path, is_datasheet)
+                    
+                    # Check if file should be processed
+                    should_process, should_remove_old = self._check_file_in_docket(file_path, metadata.md5)
+                    
+                    if not should_process:
+                        # Skip this file
+                        self.skipped_files += 1
+                        logger.info(f"Skipping file: {file_path}")
+                        continue
+                    
+                    # If we need to remove old version first
+                    if should_remove_old:
+                        existing_entry = self._docket.get_by_path(file_path)
+                        if existing_entry:
+                            logger.info(f"Removing old version from knowledge base: {file_path}")
+                            self._remove_file_from_kb(existing_entry)
+                    
+                    # Process the file
+                    logger.info(f"Pre-processing file: {file_path}")
+                    file = await self._preprocess_file(file_path, is_datasheet)
+                    self.successfully_processed += 1
+                    self._add_to_file_docket(file)
+                    
+                    # Track the file subtype as processed
+                    if file.sub_type:
+                        self._processed_file_types.add(FileSubtypes(file.sub_type).name)
+                    
+                except Exception as e:
+                    self.failed_files += 1
+                    self.errors.append({"file_path": file_path, "error": str(e)})
+                    logger.exception(e)
+                    logger.error(f"Error processing file: {file_path}")
+
+    async def run(self):
+        try:
+            logger.info(f"Starting pipeline for directory: {self.folder_path}")
+            
+            # Check if context folder has any files
+            if not self._check_context_folder_has_files():
+                self._display_empty_context_message()
+                return
+            
+            self.pipeline_run = create_pipeline_run(self._db, self._project)
+            self._context.run = self.pipeline_run
+            await self._preprocess_folder(self.folder_path, False)
+            logger.info(f"Finished pre-processing folder: {self.folder_path}")
+
+            # Run pre-processing on datasheets which were just downloaded
+            logger.info(f"Starting pre-processing of datasheets: {app_dirs.data_sheets_dir}")
+            await self._preprocess_folder(app_dirs.data_sheets_dir, True)
+            logger.info(f"Finished pre-processing of datasheets: {app_dirs.data_sheets_dir}")
+
+            logger.info(f"Preprocessing finished: {self.folder_path}")
+            logger.info(f"Running analysis step: {self.folder_path}")
+            await self._run_analysis()
+            self._db.commit()
+            self._save_file_docket()
+            report = json.dumps({
+                "total_files": self.total_files,
+                "successfully_processed": self.successfully_processed,
+                "skipped_files": self.skipped_files,
+                "failed_files": self.failed_files,
+                "errors": self.errors
+            })
+            logger.info(f"Finished pipeline")
+            logger.info(f"Report: {report}")
+            
+            # Print summary to console
+            print(f"\n{'='*70}")
+            print(f"Pipeline Execution Summary")
+            print(f"{'='*70}")
+            print(f"Total files found: {self.total_files}")
+            print(f"Successfully processed: {self.successfully_processed}")
+            print(f"Skipped (already processed): {self.skipped_files}")
+            print(f"Failed: {self.failed_files}")
+            print(f"{'='*70}\n")
+        except Exception as e:
+            logger.exception(e)
+            logger.error(f"Error in pipeline: {e}")
+            return format_error_for_llm(e)
+
+
+async def run_pipeline_with_config(project_config: ProjectConfig):
+    with Session() as db:
+        project = get_project_by_name(db, project_config.name)
+        return await PipelineRunner(db, project, project_config).run()
+
+
+async def run_pipeline(project_name: str) -> str:
+    """
+    The controller agent will call this tool to start the pipeline processing for all the files in a directory.
+    :param project_name: The name of the project
+    :return: Status of the pipeline run
+    """
+    with Session() as db:
+        project = get_project_by_name(db, project_name)
+        project_config = read_project_config_file()
+        return await PipelineRunner(db, project, project_config).run()

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)