From fb019fd4cf5e24d3e02b21d9e5942bc83ccebfaa Mon Sep 17 00:00:00 2001 From: Manish Kumar Date: Mon, 31 Aug 2026 20:20:30 -0500 Subject: [PATCH] feat: read repo list and model name from config files, defaulting to prior hardcoded behavior Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_012RttwcxBgsnXH2TwQGQsrp --- configs/index_config.yaml | 7 ++++ rag/engine.py | 35 +++++++++++++++++- requirements.txt | 1 + scripts/build_index.py | 77 +++++++++++++++++++++++++++++---------- 4 files changed, 98 insertions(+), 22 deletions(-) diff --git a/configs/index_config.yaml b/configs/index_config.yaml index e69de29..6a9cb2c 100644 --- a/configs/index_config.yaml +++ b/configs/index_config.yaml @@ -0,0 +1,7 @@ +# Ollama generation model used by rag/engine.py's ollama_generate() default +# arg and RAGEngine.stream_llm(). Both call sites read this same value. +# +# If this file is missing, empty, or llm_model isn't set, the code falls +# back to "llama3" -- the prior hardcoded value -- so behavior is unchanged +# unless this is explicitly edited. +llm_model: llama3 diff --git a/rag/engine.py b/rag/engine.py index 44194bf..a434f2f 100644 --- a/rag/engine.py +++ b/rag/engine.py @@ -4,6 +4,7 @@ import logging import requests import numpy as np +import yaml from typing import Generator, List, Dict, Any, Optional logger = logging.getLogger(__name__) @@ -33,6 +34,36 @@ def _get_cross_encoder(): OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://ollama:11434/api") EMBED_DIM = 768 +_INDEX_CONFIG_PATH = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "..", "configs", "index_config.yaml" +) +_DEFAULT_LLM_MODEL = "llama3" + + +def _load_llm_model(default: str = _DEFAULT_LLM_MODEL, path: str = _INDEX_CONFIG_PATH) -> str: + """Read `llm_model` from configs/index_config.yaml. + + Falls back to `default` if the file is missing, empty, fails to parse, + or doesn't set llm_model -- so behavior is unchanged unless the config + is explicitly edited. Both ollama_generate()'s default arg and + RAGEngine.stream_llm() read this same LLM_MODEL value. + """ + try: + with open(path, "r") as f: + data = yaml.safe_load(f) + except (OSError, yaml.YAMLError) as e: + logger.warning(f"Could not read {path} ({e}); defaulting llm_model to {default!r}") + return default + + if not isinstance(data, dict): + return default + + value = data.get("llm_model") + return str(value) if value else default + + +LLM_MODEL = _load_llm_model() + def ollama_embed(text: str, model: str = "nomic-embed-text"): # Ollama's llama-server subprocess can transiently fail CUDA context @@ -80,7 +111,7 @@ def ollama_embed(text: str, model: str = "nomic-embed-text"): return vec -def ollama_generate(prompt: str, model: str = "llama3"): +def ollama_generate(prompt: str, model: str = LLM_MODEL): res = requests.post( f"{OLLAMA_URL}/generate", json={ @@ -271,7 +302,7 @@ def stream_llm(self, query: str, context: str) -> Generator[str, None, None]: try: res = requests.post( f"{OLLAMA_URL}/generate", - json={"model": "llama3", "prompt": prompt, "stream": True}, + json={"model": LLM_MODEL, "prompt": prompt, "stream": True}, stream=True, timeout=300, ) diff --git a/requirements.txt b/requirements.txt index b10b169..b6925a7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,3 +5,4 @@ numpy==1.26.4 faiss-cpu==1.15.0 sentence-transformers==5.4.1 PyJWT==2.10.1 +PyYAML==6.0.2 diff --git a/scripts/build_index.py b/scripts/build_index.py index 7980410..64c8775 100644 --- a/scripts/build_index.py +++ b/scripts/build_index.py @@ -2,6 +2,7 @@ import os import hashlib import numpy as np +import yaml sys.path.append(os.path.abspath(".")) @@ -12,10 +13,41 @@ MIN_CHUNK_CHARS = 10 # discard overflow tails from chunker.py's hard char-slice +REPOS_CONFIG_PATH = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "..", "configs", "repos.yaml" +) + + def hash_text(text: str) -> str: return hashlib.sha256(text.encode("utf-8")).hexdigest() +def _load_repo_names_from_config(path: str = REPOS_CONFIG_PATH): + """Read the `repos:` name list from configs/repos.yaml. + + Returns the list of names on success, or None if the file is missing, + empty, fails to parse, or doesn't define a non-empty `repos:` list. In + every one of those cases the caller falls back to the hardcoded list in + build_index(), so behavior is unchanged unless the YAML is actually + populated and wired in. + """ + try: + with open(path, "r") as f: + data = yaml.safe_load(f) + except (OSError, yaml.YAMLError) as e: + print(f"⚠️ Could not read {path} ({e}); falling back to hardcoded repo list") + return None + + if not isinstance(data, dict): + return None + + names = data.get("repos") + if not names or not isinstance(names, list): + return None + + return list(names) + + def normalize_vector(vec): vec = np.array(vec, dtype=np.float32) @@ -39,28 +71,33 @@ def build_index(): # Locally: export REPO_BASE=/home/manish/Desktop/machine (or set in .env). REPO_BASE = os.environ.get("REPO_BASE", "/home/manish/Desktop/machine") - repos = [ - f"{REPO_BASE}/omnibioai", - f"{REPO_BASE}/omnibioai-rag", - f"{REPO_BASE}/omnibioai-toolserver", - f"{REPO_BASE}/omnibioai-sdk", - f"{REPO_BASE}/omnibioai-workflow-bundles", - f"{REPO_BASE}/omnibioai-control-center", - f"{REPO_BASE}/omnibioai-lims", - f"{REPO_BASE}/omnibioai-model-registry", - f"{REPO_BASE}/omnibioai-dev-docker", - f"{REPO_BASE}/omnibioai-api-gateway", - f"{REPO_BASE}/omnibioai-docs", - f"{REPO_BASE}/omnibioai-studio", - f"{REPO_BASE}/omnibioai-auth", - f"{REPO_BASE}/omnibioai-tool-runtime", - f"{REPO_BASE}/omnibioai-iam-client", - f"{REPO_BASE}/omnibioai-policy-engine", - f"{REPO_BASE}/omnibioai-security-audit", - f"{REPO_BASE}/omnibioai-security-sdk", - f"{REPO_BASE}/omnibioai-hpc-policy-engine", + # Hardcoded fallback -- kept in place (not deleted) so behavior is + # unchanged if configs/repos.yaml is missing, empty, or fails to parse. + HARDCODED_REPO_NAMES = [ + "omnibioai", + "omnibioai-rag", + "omnibioai-toolserver", + "omnibioai-sdk", + "omnibioai-workflow-bundles", + "omnibioai-control-center", + "omnibioai-lims", + "omnibioai-model-registry", + "omnibioai-dev-docker", + "omnibioai-api-gateway", + "omnibioai-docs", + "omnibioai-studio", + "omnibioai-auth", + "omnibioai-tool-runtime", + "omnibioai-iam-client", + "omnibioai-policy-engine", + "omnibioai-security-audit", + "omnibioai-security-sdk", + "omnibioai-hpc-policy-engine", ] + repo_names = _load_repo_names_from_config() or HARDCODED_REPO_NAMES + repos = [f"{REPO_BASE}/{name}" for name in repo_names] + # Fail fast: if not a single repo exists, the REPO_BASE is wrong. existing = [r for r in repos if os.path.isdir(r)] if not existing: