cnhkmcp 2.0.4__py3-none-any.whl → 2.1.1__py3-none-any.whl

cnhkmcp/untracked/AI/321/206/320/261/320/234/321/211/320/255/320/262/321/206/320/237/320/242/321/204/342/225/227/342/225/242/process_knowledge_base.py

@@ -0,0 +1,280 @@
+import os
+import sys
+import pandas as pd
+import json
+import re
+from pathlib import Path
+from typing import List
+
+# Add get_knowledgeBase_tool to path
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+TOOL_DIR = os.path.join(SCRIPT_DIR, "get_knowledgeBase_tool")
+if TOOL_DIR not in sys.path:
+    sys.path.insert(0, TOOL_DIR)
+
+# Import from tool directory
+sys.path.insert(0, TOOL_DIR)
+import ace_lib
+from fetch_all_operators import fetch_operators, prompt_credentials
+from fetch_all_datasets import (
+    fetch_all_combinations,
+    fetch_datasets_for_combo,
+    merge_and_deduplicate,
+)
+from fetch_all_documentation import (
+    fetch_tutorials,
+    fetch_tutorial_pages,
+    fetch_page,
+    _extract_page_id,
+)
+
+
+def ensure_knowledge_dir():
+    """Ensure knowledge directory exists"""
+    knowledge_dir = os.path.join(SCRIPT_DIR, "knowledge")
+    os.makedirs(knowledge_dir, exist_ok=True)
+    return knowledge_dir
+
+
+def to_jsonable(value):
+    """Convert values to JSON-serializable, handling NaN and nested structures."""
+    try:
+        if isinstance(value, float) and pd.isna(value):
+            return None
+    except TypeError:
+        pass
+
+    if isinstance(value, list):
+        return [to_jsonable(v) for v in value if not (isinstance(v, float) and pd.isna(v))]
+    if isinstance(value, dict):
+        return {k: to_jsonable(v) for k, v in value.items()}
+    if isinstance(value, (str, int, float, bool)) or value is None:
+        return value
+    return str(value)
+
+
+def safe_filename(name: str, suffix: str = "") -> str:
+    base = re.sub(r"[^A-Za-z0-9._-]+", "_", str(name)).strip("_") or "doc"
+    base = base[:80]
+    return f"{base}{suffix}"
+
+
+def process_operators(session: ace_lib.SingleSession, knowledge_dir: str):
+    """
+    Process operators and save as JSON files
+    
+    Args:
+        session: Authenticated BRAIN session
+        knowledge_dir: Directory to save JSON files
+    """
+    print("\n=== Processing Operators ===")
+    
+    # Fetch operators data
+    print("Fetching operators...")
+    operators_df = fetch_operators(session)
+    
+    if operators_df.empty:
+        print("No operators found!")
+        return
+    
+    print(f"Found {len(operators_df)} operator entries")
+    
+    # Get unique categories
+    categories = sorted(operators_df['category'].dropna().unique())
+    
+    for category in categories:
+        category_data = operators_df[operators_df['category'] == category].copy()
+        
+        # Create JSON file for this category
+        filename = f"{category.replace(' ', '_').lower()}_operators.json"
+        filepath = os.path.join(knowledge_dir, filename)
+        
+        print(f"Processing category: {category}")
+        
+        # Convert to list of dicts
+        category_list = []
+        for idx, row in category_data.iterrows():
+            operator_dict = {}
+            for col in row.index:
+                value = row[col]
+                operator_dict[col] = to_jsonable(value)
+            category_list.append(operator_dict)
+        
+        # Save category JSON
+        with open(filepath, 'w', encoding='utf-8') as f:
+            json.dump(category_list, f, ensure_ascii=False, indent=2)
+        
+        print(f"✓ Created {filename} with {len(category_list)} operators")
+
+
+def process_datasets(session: ace_lib.SingleSession, dataset_dir: str):
+    """Fetch datasets and save one JSON per region."""
+    print("=== Processing Datasets ===")
+
+    print("Fetching valid instrument/region/delay/universe combinations...")
+    options_df = fetch_all_combinations(session)
+    if options_df is None or options_df.empty:
+        print("No simulation options fetched; aborting dataset fetch.")
+        return
+
+    all_datasets: list[pd.DataFrame] = []
+    combo_idx = 0
+
+    for _, row in options_df.iterrows():
+        instrument_type = row.get("InstrumentType")
+        region = row.get("Region")
+        delay = row.get("Delay")
+        universes = row.get("Universe") or []
+
+        for universe in universes:
+            combo_idx += 1
+            print(f"[{combo_idx}] {instrument_type} / {region} / D{delay} / {universe}")
+            try:
+                df = fetch_datasets_for_combo(session, instrument_type, region, delay, universe)
+                print(f"    -> {len(df)} rows")
+                all_datasets.append(df)
+            except Exception as exc:
+                print(f"    -> Failed: {exc}")
+
+    if not all_datasets:
+        print("No datasets fetched; nothing to save.")
+        return
+
+    combined_df = pd.concat([df for df in all_datasets if not df.empty], ignore_index=True)
+    if combined_df.empty:
+        print("No datasets fetched; nothing to save.")
+        return
+
+    regions = sorted(combined_df["param_region"].dropna().unique())
+    print(f"Found regions: {', '.join(regions)}")
+
+    for region in regions:
+        region_df = combined_df[combined_df["param_region"] == region]
+        region_unique = merge_and_deduplicate([region_df])
+
+        region_list = []
+        for _, row in region_unique.iterrows():
+            record = {col: to_jsonable(row[col]) for col in row.index}
+            region_list.append(record)
+
+        filename = f"{region.replace(' ', '_').lower()}_datasets.json"
+        filepath = os.path.join(dataset_dir, filename)
+        with open(filepath, "w", encoding="utf-8") as f:
+            json.dump(region_list, f, ensure_ascii=False, indent=2)
+
+        print(f"✓ Created {filename} with {len(region_list)} datasets")
+
+
+def process_documentation(session: ace_lib.SingleSession, knowledge_dir: str):
+    """Fetch tutorials and pages, save one JSON per page."""
+    print("=== Processing Documentation ===")
+
+    tutorials = fetch_tutorials(session)
+    if not tutorials:
+        print("No tutorials fetched; skipping documentation.")
+        return
+
+    print(f"Fetched {len(tutorials)} tutorials")
+
+    page_count = 0
+    seen_pages = set()
+    
+    for idx, tutorial in enumerate(tutorials, start=1):
+        tutorial_id = _extract_page_id(tutorial) or f"tutorial_{idx}"
+        tutorial_title = tutorial.get("title") or tutorial_id
+
+        page_candidates = []
+        if isinstance(tutorial.get("pages"), list):
+            page_candidates.extend(tutorial["pages"])
+        if tutorial_id:
+            try:
+                page_candidates.extend(fetch_tutorial_pages(session, tutorial_id))
+            except Exception as exc:
+                print(f"[{idx:03d}] failed to fetch pages for {tutorial_id}: {exc}")
+
+        if not page_candidates and tutorial_id:
+            page_candidates.append({"id": tutorial_id, "title": tutorial_title})
+
+        for page_entry in page_candidates:
+            page_id = _extract_page_id(page_entry)
+            if not page_id or page_id in seen_pages:
+                continue
+            seen_pages.add(page_id)
+            
+            try:
+                page = fetch_page(session, page_id)
+            except Exception as exc:
+                print(f"[{idx:03d}] page {page_id} failed: {exc}")
+                continue
+            
+            page_count += 1
+            page_title = page.get("title") or page_entry.get("title") or page_id
+            
+            # Save each page as individual JSON
+            filename = safe_filename(f"{idx:03d}_{page_title}", "_documentation.json")
+            filepath = os.path.join(knowledge_dir, filename)
+            
+            with open(filepath, "w", encoding="utf-8") as f:
+                json.dump(to_jsonable(page), f, ensure_ascii=False, indent=2)
+            
+            print(f"[{idx:03d}] ✓ Created {filename}")
+    
+    print(f"✓ Total: {page_count} documentation pages saved")
+
+
+def main():
+    print("=== BRAIN Knowledge Base Processor ===")
+    print("Starting operator processing...\n")
+    
+    # Get credentials
+    email, password = prompt_credentials()
+    ace_lib.get_credentials = lambda: (email, password)
+    
+    print("Logging in to BRAIN platform...")
+    try:
+        session = ace_lib.start_session()
+        print("✓ Login successful\n")
+    except Exception as exc:
+        print(f"✗ Login failed: {exc}")
+        return
+    
+    # Ensure knowledge directory exists
+    knowledge_dir = ensure_knowledge_dir()
+    dataset_dir = knowledge_dir  # Save datasets directly under knowledge
+    print(f"Knowledge directory: {knowledge_dir}\n")
+
+    # Process documentation (tutorials/pages)
+    print("\nStarting documentation processing...\n")
+    try:
+        process_documentation(session, knowledge_dir)
+    except Exception as exc:
+        print(f"✗ Failed to process documentation: {exc}")
+        import traceback
+        traceback.print_exc()
+        return
+
+    # Process operators
+    try:
+        process_operators(session, knowledge_dir)
+    except Exception as exc:
+        print(f"✗ Failed to process operators: {exc}")
+        import traceback
+        traceback.print_exc()
+        return
+
+    # Process datasets by region
+    print("\nStarting dataset processing...\n")
+    try:
+        process_datasets(session, dataset_dir)
+    except Exception as exc:
+        print(f"✗ Failed to process datasets: {exc}")
+        import traceback
+        traceback.print_exc()
+        return
+
+
+    print("\n=== Processing Complete ===")
+
+
+if __name__ == "__main__":
+    main()

cnhkmcp/untracked/AI/321/206/320/261/320/234/321/211/320/255/320/262/321/206/320/237/320/242/321/204/342/225/227/342/225/242/rag_engine.py

@@ -0,0 +1,356 @@
+import os
+import json
+import shutil
+import chromadb
+from fastembed import TextEmbedding
+from watchdog.observers import Observer
+from watchdog.events import FileSystemEventHandler
+import threading
+
+
+PREFERRED_MODELS = [
+    "jinaai/jina-embeddings-v2-base-zh",              # 中英混合友好，~0.64GB
+    "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",  # 多语 ~50 语种
+    "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",  # 多语轻量版
+    "intfloat/multilingual-e5-large",                 # 多语更强，体积约 2.2GB
+]
+
+# Final chosen model will be detected at runtime from supported list
+MODEL_NAME = None
+COLLECTION_NAME = "brain_kb_v5"
+
+# Optional imports for different file types
+try:
+    from pypdf import PdfReader
+except ImportError:
+    PdfReader = None
+
+try:
+    from docx import Document
+except ImportError:
+    Document = None
+
+class KnowledgeBase:
+    def __init__(self, kb_path="knowledge", db_path="vector_db"):
+        self.kb_path = os.path.abspath(kb_path)
+        self.db_path = os.path.abspath(db_path)
+        self.meta_path = os.path.join(self.db_path, "_meta.json")
+        self.manifest_path = os.path.join(self.db_path, "_manifest.json")
+        self._collection_reset_guard = False
+        
+        if not os.path.exists(self.kb_path):
+            os.makedirs(self.kb_path)
+        
+        # Initialize Embedding Model (BAAI/bge-small-zh-v1.5 is ~100MB)
+        # This will load from cache if already downloaded
+        # Pick the first available model from the preferred list
+        _supported_raw = TextEmbedding.list_supported_models()
+        supported = set()
+        for item in _supported_raw:
+            if isinstance(item, dict) and "model" in item:
+                supported.add(item["model"])
+            elif isinstance(item, str):
+                supported.add(item)
+        chosen = None
+        for name in PREFERRED_MODELS:
+            if name in supported:
+                chosen = name
+                break
+        if not chosen:
+            raise RuntimeError(
+                "No preferred embedding models are supported by fastembed. "
+                "Please check available models via TextEmbedding.list_supported_models()."
+            )
+
+        print(f"Loading Knowledge Base Embedding Model: {chosen} (may take some time on first run)...")
+        try:
+            self.model = TextEmbedding(model_name=chosen)
+            print("Embedding Model loaded successfully.")
+        except Exception as e:
+            print(f"Error loading embedding model: {e}")
+            raise
+
+        # Store chosen model name for reference
+        global MODEL_NAME
+        MODEL_NAME = chosen
+
+        # Cache embedding dimension (detects library/model changes that corrupt existing indexes)
+        self.embed_dim = self._get_embedding_dim()
+        self.chroma_version = getattr(chromadb, "__version__", "unknown")
+
+        # If the stored index was built with a different model/dimension/chromadb version, wipe it
+        self._maybe_reset_for_incompatibility(chosen, self.embed_dim, self.chroma_version)
+        
+        # Initialize Vector DB
+        self._init_collection()
+        
+        # Initial sync
+        self.sync_knowledge()
+        
+        # Start Watcher
+        self.start_watcher()
+
+    def _init_collection(self, recreate: bool = False):
+        """(Re)initialize Chroma client/collection. If recreate=True, wipe on-disk index."""
+        if recreate and os.path.exists(self.db_path):
+            shutil.rmtree(self.db_path, ignore_errors=True)
+        try:
+            self.client = chromadb.PersistentClient(path=self.db_path)
+            self.collection = self.client.get_or_create_collection(
+                name=COLLECTION_NAME,
+                metadata={"hnsw:space": "cosine"}
+            )
+        except Exception as exc:
+            # If collection load itself fails, wipe and retry once to clear corrupted segments
+            if not recreate:
+                shutil.rmtree(self.db_path, ignore_errors=True)
+                return self._init_collection(recreate=True)
+            raise
+
+        # Persist metadata about the embedding model used to build this index
+        try:
+            os.makedirs(self.db_path, exist_ok=True)
+            with open(self.meta_path, "w", encoding="utf-8") as f:
+                json.dump({
+                    "model": MODEL_NAME,
+                    "embed_dim": self.embed_dim,
+                    "chroma_version": self.chroma_version,
+                }, f)
+        except Exception:
+            pass  # Metadata failure should not block runtime
+
+    def _maybe_reset_for_incompatibility(self, chosen_model: str, embed_dim: int, chroma_version: str):
+        """If existing index meta differs (model/dimension/chromadb), wipe it."""
+        if not os.path.exists(self.db_path):
+            return
+        try:
+            with open(self.meta_path, "r", encoding="utf-8") as f:
+                meta = json.load(f)
+            prev_model = meta.get("model")
+            prev_dim = meta.get("embed_dim")
+            prev_chroma = meta.get("chroma_version")
+            if prev_model != chosen_model or prev_dim != embed_dim or prev_chroma != chroma_version:
+                shutil.rmtree(self.db_path, ignore_errors=True)
+        except Exception:
+            # If meta cannot be read, assume stale/corrupted and rebuild
+            shutil.rmtree(self.db_path, ignore_errors=True)
+
+    def _get_embedding_dim(self) -> int:
+        for vec in self.model.embed(["dimension_probe"]):
+            try:
+                return len(vec)
+            except Exception:
+                return len(list(vec))
+        raise RuntimeError("Failed to determine embedding dimension")
+
+    def sync_knowledge(self, allow_reset: bool = True):
+        """Scans the knowledge folder and updates the vector database."""
+        print("Syncing knowledge base...")
+        manifest = self._load_manifest()
+        updated_manifest = {}
+        supported_extensions = (".txt", ".md", ".pdf", ".docx", ".json")
+        current_files = []
+        for filename in os.listdir(self.kb_path):
+            file_path = os.path.join(self.kb_path, filename)
+            if os.path.isfile(file_path) and filename.lower().endswith(supported_extensions):
+                current_files.append(filename)
+                mtime = os.path.getmtime(file_path)
+                size = os.path.getsize(file_path)
+                prev_meta = manifest.get(filename)
+                # Skip unchanged files
+                if prev_meta and prev_meta.get("mtime") == mtime and prev_meta.get("size") == size:
+                    updated_manifest[filename] = prev_meta
+                    continue
+                try:
+                    content = self._extract_text(file_path)
+                    if content:
+                        # Sliding window chunking on original text
+                        chunk_size = 400
+                        overlap = 80
+                        original_chunks = []
+                        for i in range(0, len(content), chunk_size - overlap):
+                            chunk = content[i:i + chunk_size].strip()
+                            if chunk:
+                                original_chunks.append(chunk)
+                        
+                        if original_chunks:
+                            # Normalize for embedding generation only (not for storage)
+                            normalized_chunks = [c.lower().replace('_', ' ') for c in original_chunks]
+                            
+                            ids = [f"{filename}_{i}" for i in range(len(original_chunks))]
+                            metadatas = [{"source": filename, "chunk": i} for i in range(len(original_chunks))]
+                            
+                            # Compute embeddings from normalized text
+                            embeddings = []
+                            for v in self.model.embed(normalized_chunks):
+                                try:
+                                    embeddings.append(v.tolist())
+                                except Exception:
+                                    embeddings.append(list(v))
+                            
+                            # Store ORIGINAL text (not normalized) so users see the real content
+                            self.collection.upsert(
+                                documents=original_chunks,
+                                ids=ids,
+                                metadatas=metadatas,
+                                embeddings=embeddings
+                            )
+                            print(f"  ✓ Indexed {filename}: {len(original_chunks)} chunks")
+                            updated_manifest[filename] = {"mtime": mtime, "size": size}
+                except Exception as e:
+                    err_msg = str(e)
+                    print(f"Error processing {filename}: {err_msg}")
+                    # Auto-recover if HNSW/compaction/index errors occur
+                    if allow_reset and any(x in err_msg.lower() for x in ["hnsw", "compaction", "segment reader"]):
+                        if not self._collection_reset_guard:
+                            print("Detected index corruption. Rebuilding vector_db and retrying sync once...")
+                            self._collection_reset_guard = True
+                            self._init_collection(recreate=True)
+                            return self.sync_knowledge(allow_reset=False)
+        # Remove deleted files from the index
+        deleted_files = set(manifest.keys()) - set(current_files)
+        for filename in deleted_files:
+            try:
+                self.collection.delete(where={"source": filename})
+                print(f"  ✓ Removed deleted file from index: {filename}")
+            except Exception as e:
+                print(f"  ! Failed to remove {filename}: {e}")
+        # Persist manifest
+        self._save_manifest(updated_manifest)
+        print("Knowledge base sync complete.")
+
+    def _extract_text(self, file_path):
+        ext = os.path.splitext(file_path)[1].lower()
+        if ext == ".txt":
+            with open(file_path, 'r', encoding='utf-8') as f:
+                return f.read()
+        elif ext == ".md":
+            # Treat Markdown as plain text for retrieval
+            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
+                return f.read()
+        elif ext == ".pdf":
+            if PdfReader:
+                reader = PdfReader(file_path)
+                text = ""
+                for page in reader.pages:
+                    text += page.extract_text() + "\n"
+                return text
+            else:
+                print("pypdf not installed, skipping PDF.")
+        elif ext == ".docx":
+            if Document:
+                doc = Document(file_path)
+                return "\n".join([para.text for para in doc.paragraphs])
+            else:
+                print("python-docx not installed, skipping Word.")
+        elif ext == ".json":
+            with open(file_path, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+                return json.dumps(data, ensure_ascii=False, indent=2)
+        return None
+
+    def query(self, text, top_k=5, distance_threshold=0.8):
+        """Retrieves relevant snippets from the knowledge base.
+
+        Uses cosine distance (lower is better). A result is treated as a hit only
+        when best_distance <= distance_threshold.
+        Returns:
+            dict: {"hit": bool, "context": str, "hits": [{source, chunk, distance, text}, ...]}
+        """
+        try:
+            # Normalize query same as indexed content
+            normalized_text = text.lower().replace('_', ' ')
+            
+            q_vec = None
+            for v in self.model.embed([normalized_text]):
+                try:
+                    q_vec = v.tolist()
+                except Exception:
+                    q_vec = list(v)
+                break
+            if q_vec is None:
+                return {"hit": False, "context": "", "hits": []}
+
+            results = self.collection.query(
+                query_embeddings=[q_vec],
+                n_results=top_k,
+                include=["documents", "metadatas", "distances"]
+            )
+
+            docs = (results or {}).get("documents") or []
+            metas = (results or {}).get("metadatas") or []
+            dists = (results or {}).get("distances") or []
+
+            if not docs or not docs[0]:
+                print("[KB Query] No results returned from collection")
+                return {"hit": False, "context": "", "hits": []}
+
+            docs0 = docs[0]
+            metas0 = metas[0] if metas and metas[0] else [{} for _ in docs0]
+            dists0 = dists[0] if dists and dists[0] else [None for _ in docs0]
+
+            hits = []
+            for doc_text, meta, dist in zip(docs0, metas0, dists0):
+                hits.append({
+                    "source": (meta or {}).get("source", ""),
+                    "chunk": (meta or {}).get("chunk", None),
+                    "distance": dist,
+                    "text": doc_text,
+                })
+
+            best = hits[0].get("distance")
+            is_hit = (best is not None) and (best <= distance_threshold)
+            
+            # Debug log
+            best_str = f"{best:.4f}" if best is not None else "N/A"
+            print(f"[KB Query] '{text[:50]}...' -> best_dist={best_str}, threshold={distance_threshold}, hit={is_hit}")
+            if hits:
+                top3_dists = [f"{h['distance']:.4f}" if h['distance'] is not None else "N/A" for h in hits[:3]]
+                print(f"[KB Query] Top 3 distances: {top3_dists}")
+            
+            context = "\n---\n".join([h["text"] for h in hits]) if is_hit else ""
+            return {"hit": is_hit, "context": context, "hits": hits}
+        except Exception as e:
+            print(f"Query error: {e}")
+            import traceback
+            traceback.print_exc()
+            return {"hit": False, "context": "", "hits": []}
+
+    def start_watcher(self):
+        event_handler = KBHandler(self)
+        self.observer = Observer()
+        self.observer.schedule(event_handler, self.kb_path, recursive=False)
+        self.observer.start()
+
+    def _load_manifest(self):
+        if not os.path.exists(self.manifest_path):
+            return {}
+        try:
+            with open(self.manifest_path, "r", encoding="utf-8") as f:
+                return json.load(f)
+        except Exception:
+            return {}
+
+    def _save_manifest(self, data):
+        try:
+            os.makedirs(self.db_path, exist_ok=True)
+            with open(self.manifest_path, "w", encoding="utf-8") as f:
+                json.dump(data, f, ensure_ascii=False, indent=2)
+        except Exception as e:
+            print(f"  ! Failed to save manifest: {e}")
+
+class KBHandler(FileSystemEventHandler):
+    def __init__(self, kb_instance):
+        self.kb = kb_instance
+        self.supported_extensions = (".txt", ".md", ".pdf", ".docx", ".json")
+
+    def on_modified(self, event):
+        if not event.is_directory and event.src_path.lower().endswith(self.supported_extensions):
+            print(f"File modified: {event.src_path}. Re-syncing...")
+            threading.Thread(target=self.kb.sync_knowledge).start()
+
+    def on_created(self, event):
+        if not event.is_directory and event.src_path.lower().endswith(self.supported_extensions):
+            print(f"File created: {event.src_path}. Syncing...")
+            threading.Thread(target=self.kb.sync_knowledge).start()
+

cnhkmcp/untracked/AI/321/206/320/261/320/234/321/211/320/255/320/262/321/206/320/237/320/242/321/204/342/225/227/342/225/242/requirements.txt

@@ -0,0 +1,7 @@
+openai
+Pillow
+fastembed>=0.3.4
+chromadb>=0.5.0
+watchdog
+pypdf
+python-docx

cnhkmcp/untracked/AI/321/206/320/261/320/234/321/211/320/255/320/262/321/206/320/237/320/242/321/204/342/225/227/342/225/242/run.bat

@@ -0,0 +1,3 @@
+@echo off
+python main.py
+pause