offloaded main agent to private repo

app.py

@@ -1,196 +1,38 @@
 import os
-import shutil
 from pathlib import Path
 import gradio as gr
+from huggingface_hub import snapshot_download, hf_hub_download
+import importlib.util
 
-from huggingface_hub import snapshot_download
-from langchain_community.vectorstores import FAISS
-from langchain_huggingface import HuggingFaceEmbeddings
+# This is the path to your private dataset repository on Hugging Face Hub
+PRIVATE_DATASET_ID = os.getenv("PRIVATE_DATASET_ID")
+HF_TOKEN = os.getenv("HF_TOKEN")
+INDEX_SUBDIR = os.getenv("INDEX_SUBDIR", ".")
 
-# ---- Your class from the prompt (unchanged except 1 tweak: import path safety) ----
-import re
-from pathlib import Path as _Path
+# --- Core Logic Download and Import ---
 try:
-    from openai import OpenAI
-except ImportError:
-    OpenAI = None
-
-class VerilogAgent:
-    """
-    A self-contained agent for generating Verilog code using API-based LLMs (e.g., GPT)
-    and a RAG pipeline.
-    """
-    def __init__(self, model_id, embedding_id, faiss_index_path, api_key):
-        self.model_id = model_id
-        self.embedding_id = embedding_id
-        self.faiss_index_path = faiss_index_path
-        self.api_key = api_key
-        
-        print(f"[INFO] Initializing VerilogAgent for model: {self.model_id}")
-        self._load_dependencies()
-
-    def _load_dependencies(self):
-        print(f"[INFO] Loading embedding model '{self.embedding_id}'...")
-        embedding = HuggingFaceEmbeddings(model_name=self.embedding_id)
-
-        print(f"[INFO] Loading FAISS vector store from '{self.faiss_index_path}'...")
-        if not _Path(self.faiss_index_path).exists():
-            raise FileNotFoundError(f"FAISS index directory not found at {self.faiss_index_path}.")
-        self.vectorstore = FAISS.load_local(self.faiss_index_path, embedding, allow_dangerous_deserialization=True)
-        
-        if not OpenAI:
-            raise ImportError("OpenAI library is not installed. Please add 'openai' to requirements.txt.")
-        if not self.api_key:
-            raise ValueError("OpenAI API key is required.")
-        self.client = OpenAI(api_key=self.api_key)
-        print("[INFO] OpenAI client initialized.")
-        print("[INFO] VerilogAgent initialized successfully.")
-
-    def _build_prompt_messages_rag(self, query: str, docs: list = None) -> list:
-        context_section = ""
-        if docs:
-            context = "\n\n".join([doc.page_content for doc in docs])
-            context_section = f"""
-CONTEXT EXAMPLES:
-```verilog
-{context}
-```"""
-
-        system_prompt = """You are an expert Verilog code generation assistant.
-TASK: Generate fully implemented, syntactically correct Verilog code in response to the user request.
-
-INSTRUCTIONS:
-1. Analyze the user request and the provided context examples to determine the required modules and logic.
-2. The context provides examples of valid, complete Verilog modules.
-3. Implement all required modules. Every `module ... endmodule` block must be complete.
-4. Do not leave logic empty or use placeholders like `// your code here`.
-5. Your entire response MUST be only the Verilog code, wrapped in a single `verilog` markdown block. Do not include any natural language explanations.
-6. Use only Verilog-2005 syntax. Do not use SystemVerilog constructs (e.g., `logic`, `always_ff`).
-7. Ensure all identifiers are declared before use and vector ranges are ordered [MSB:LSB].
-"""
-        user_prompt = f'''TASK: Generate a fully implemented, syntactically correct Verilog module named 'TopModule'. This name is a strict requirement.
-{context_section}
-
-USER REQUEST:
-"""{query}"""
-
-OUTPUT:
-Generate the complete Verilog code.
-'''
-        return [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": user_prompt}
-        ]
-
-    def _build_prompt_messages_baseline(self, query: str, docs: list = None) -> list:
-        context_section = ""
-        if docs:
-            context = "\n\n".join([doc.page_content for doc in docs])
-            context_section = f"""
-CONTEXT EXAMPLES:
-```verilog
-{context}
-```"""
-        system_prompt = "You are an expert Verilog code generation assistant."
-        user_prompt = f'''TASK: Generate a fully implemented, syntactically correct Verilog module named 'TopModule'. This name is a strict requirement.
-{context_section}
-
-USER REQUEST:
-"""{query}"""
-
-OUTPUT:
-Generate the complete Verilog code.
-'''
-        return [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": user_prompt}
-        ]
-
-    def _extract_verilog_code(self, text: str) -> str:
-        verilog_pattern = re.compile(r"```(?:verilog\s*)?(.*?)\s*```", re.DOTALL)
-        match = verilog_pattern.search(text)
-        if match:
-            return match.group(1).strip()
-        module_pattern = re.compile(r"(module.*?endmodule)", re.DOTALL)
-        match = module_pattern.search(text)
-        if match:
-            return match.group(1).strip()
-        return text.strip()
-
-    def _call_api(self, messages: list, generation_params: dict) -> str:
-        try:
-            api_params = {
-                "model": self.model_id,
-                "messages": messages,
-                "max_tokens": generation_params.get("max_new_tokens"),
-                "temperature": generation_params.get("temperature"),
-                "top_p": generation_params.get("top_p")
-            }
-            if "gpt-5" in self.model_id and "verbosity" in generation_params:
-                api_params["verbosity"] = generation_params["verbosity"]
-
-            completion = self.client.chat.completions.create(**api_params)
-            return self._extract_verilog_code(completion.choices[0].message.content)
-        except Exception as e:
-            print(f"[ERROR] Code generation failed for model {self.model_id}: {e}")
-            return f"// ERROR: Generation failed. Details: {e}"
-
-    def generate_with_context(self, spec: str, docs_with_scores: list, generation_params: dict) -> str:
-        relevant_docs = [doc for doc, score in docs_with_scores]
-        messages = self._build_prompt_messages_rag(spec, relevant_docs)
-        return self._call_api(messages, generation_params)
-
-    def generate_baseline(self, spec: str, generation_params: dict) -> str:
-        messages = self._build_prompt_messages_baseline(spec, docs=[])
-        return self._call_api(messages, generation_params)
-
-
-# --------------------------- Space wiring below ---------------------------
-
-CACHE_DIR = Path("./faiss_index")  # Stores index in repo's ephemeral environment
-CACHE_DIR.mkdir(parents=True, exist_ok=True)
-
-# Env vars you’ll set in the Space “Settings → Repository secrets”
-HF_TOKEN = os.getenv("HF_TOKEN")  # personal access token with read permission
-PRIVATE_DATASET_ID = os.getenv("PRIVATE_DATASET_ID")  # e.g. "yourname/VerilogDB_faiss"
-INDEX_SUBDIR = os.getenv("INDEX_SUBDIR", ".")  # since your files are at repo root
-EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
-
-def ensure_index_downloaded() -> Path:
-    """
-    Downloads your private dataset (FAISS index + artifacts) once per container.
-    Avoids committing large binaries to the public Space repo.
-    """
-    target = CACHE_DIR / INDEX_SUBDIR
-    if target.exists() and any(target.iterdir()):
-        print(f"[INFO] Using cached FAISS index at {target}")
-        return target
-
-    if not HF_TOKEN:
-        raise RuntimeError("Missing HF_TOKEN secret. Add it in the Space settings.")
-    if not PRIVATE_DATASET_ID:
-        raise RuntimeError("Missing PRIVATE_DATASET_ID secret (e.g., 'user/VerilogDB_faiss').")
-
-    print(f"[INFO] Downloading private dataset: {PRIVATE_DATASET_ID}")
-    snapshot_path = snapshot_download(
+    # First, download the core agent code from the private repo
+    AGENT_CODE_PATH = hf_hub_download(
         repo_id=PRIVATE_DATASET_ID,
-        repo_type="dataset",
-        token=HF_TOKEN,
-        local_dir=str(CACHE_DIR),
-        local_dir_use_symlinks=False,
+        filename="deepv_core.py",  # The file containing your VerilogAgent class
+        repo_type="dataset", # Ensure this is 'dataset' to match your repo type
+        token=HF_TOKEN
     )
+    # Dynamically load the agent module from the downloaded file
+    spec = importlib.util.spec_from_file_location("deepv_core_module", AGENT_CODE_PATH)
+    agent_module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(agent_module)
+
+    # Now you can access the functions and classes from the private module
+    VerilogAgent = agent_module.VerilogAgent
+    run_generation = agent_module.run_generation
+    get_vectorstore_path = agent_module.get_vectorstore_path
+    ensure_index_downloaded = agent_module.ensure_index_downloaded
     
-    # Walk the downloaded directory to find the index files
-    downloaded_dir = Path(snapshot_path)
-    for root, dirs, files in os.walk(downloaded_dir):
-        if 'index.faiss' in files and 'index.pkl' in files:
-            index_path = Path(root)
-            print(f"[INFO] Found index files at {index_path}. Using this path.")
-            return index_path
-
-    # If the files are not found, raise an error
-    raise FileNotFoundError("FAISS index files (index.faiss and index.pkl) not found in the downloaded dataset.")
-
+except Exception as e:
+    # Handle the error gracefully if the private repo can't be accessed
+    def show_error(*args):
+        return f"// ERROR: Failed to load core agent code. Check your Hugging Face token and private dataset configuration. Details: {e}", "", []
 
 # Keep a lightweight global cache so we don’t reload embeddings on every click
 _VECTORSTORE_PATH = None
@@ -201,58 +43,12 @@ def get_vectorstore_path() -> Path:
         _VECTORSTORE_PATH = ensure_index_downloaded()
     return _VECTORSTORE_PATH
 
-def run_generation(spec, use_rag, top_k, model_choice, api_key, temperature, top_p, max_new_tokens):
-    if not spec or not api_key:
-        return "// Please provide a design specification and your API key.", "", []
-
-    # Prepare agent
-    try:
-        faiss_path = get_vectorstore_path()
-        agent = VerilogAgent(
-            model_id=model_choice,
-            embedding_id=EMBEDDING_MODEL,
-            faiss_index_path=str(faiss_path),
-            api_key=api_key.strip()
-        )
-    except Exception as e:
-        return f"// Initialization error: {e}", "", []
-
-    # Retrieval (if enabled)
-    docs_with_scores = []
-    retrieved_preview = []
-    retrieved_raw_formatted = [] # New list to hold formatted data
-    if use_rag:
-        try:
-            docs_with_scores = agent.vectorstore.similarity_search_with_score(spec, k=top_k)
-            for doc, score in docs_with_scores:
-                src = doc.metadata.get("source_file", doc.metadata.get("module", "unknown"))
-                retrieved_preview.append(f"{src}  |  score={score:.4f}")
-                # Add the page content to the new list, formatted as a tuple
-                retrieved_raw_formatted.append((doc.page_content, None))
-        except Exception as e:
-            return f"// Retrieval error: {e}", "", []
-
-    # Call model
-    gen_params = {
-        "temperature": float(temperature),
-        "top_p": float(top_p),
-        "max_new_tokens": int(max_new_tokens),
-    }
-
-    if use_rag:
-        code = agent.generate_with_context(spec, docs_with_scores, gen_params)
-    else:
-        code = agent.generate_baseline(spec, gen_params)
-
-    # Return the new formatted list for the HighlightedText component
-    return code.strip(), ("\n".join(retrieved_preview) if retrieved_preview else ""), retrieved_raw_formatted
-
+# --- Gradio UI setup below ---
 with gr.Blocks(title="DeepV for RTL (Model-Agnostic)") as demo:
     gr.Markdown("## DeepV for RTL Code Generation — Model-Agnostic (Bring Your Own API Key)")
 
     with gr.Row():
         with gr.Column(scale=2):
-            # Moved model choice and API key to the top of the left column
             with gr.Row():
                 model_choice = gr.Dropdown(
                     choices=[
@@ -260,13 +56,18 @@ with gr.Blocks(title="DeepV for RTL (Model-Agnostic)") as demo:
                         "gpt-4o-mini",
                         "gpt-4.1",
                         "gpt-5-chat-latest",
-                        "gpt-5-mini"
                     ],
                     value="gpt-4o",
                     label="Model"
                 )
                 api_key = gr.Textbox(label="OpenAI API Key", type="password", placeholder="sk-...")
 
+            gr.Markdown(
+                """
+                **Note:** Your API key is used for the current session only and is not saved or stored.
+                """
+            )
+
             spec = gr.Textbox(
                 label="Design Specification (natural language or I/O contract)",
                 placeholder="e.g., 8-bit UART transmitter with baud rate generator ...",
@@ -306,4 +107,11 @@ with gr.Blocks(title="DeepV for RTL (Model-Agnostic)") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    if 'agent_module' in locals():
+        demo.launch()
+    else:
+        with gr.Blocks() as error_demo:
+            gr.Markdown("# Initialization Error")
+            gr.Markdown(f"An error occurred while loading the application code. Please check your configuration.")
+            gr.Textbox(label="Error Details", value=str(e), lines=5)
+        error_demo.launch()