Refactor environment configuration and remove Vertex AI dependencies

.env.example

@@ -4,8 +4,6 @@
 
 # ── General ───────────────────────────────────────────────────────────────────
 MEDIC_ENV=local                    # local | kaggle | production
-MEDIC_DEFAULT_BACKEND=local        # local | vertex
-MEDIC_USE_VERTEX=false
 MEDIC_QUANTIZATION=4bit            # none | 4bit
 
 # ── Local HuggingFace Models ──────────────────────────────────────────────────

README.md

@@ -47,8 +47,6 @@ Patient form ──► Agent 1: Intake Historian  ──► (no lab) ───
 - HuggingFace account with access granted to:
   - [MedGemma](https://huggingface.co/google/medgemma-4b-it)
   - [TxGemma](https://huggingface.co/google/txgemma-2b-predict)
-- **For cloud deployment:** Google Cloud project with Vertex AI enabled
-
 ---
 
 ## Setup
@@ -68,9 +66,6 @@ cp .env.example .env
 Edit `.env`. Minimum required settings:
 
 ```bash
-# Choose your backend
-MEDIC_DEFAULT_BACKEND=local        # local | vertex
-
 # Local model IDs (HuggingFace)
 MEDIC_LOCAL_MEDGEMMA_4B_MODEL=google/medgemma-4b-it
 MEDIC_LOCAL_MEDGEMMA_27B_MODEL=google/medgemma-4b-it   # use 4B as fallback if <24 GB VRAM
@@ -78,15 +73,6 @@ MEDIC_LOCAL_TXGEMMA_9B_MODEL=google/txgemma-2b-predict
 MEDIC_LOCAL_TXGEMMA_2B_MODEL=google/txgemma-2b-predict
 ```
 
-For Vertex AI instead:
-
-```bash
-MEDIC_DEFAULT_BACKEND=vertex
-MEDIC_USE_VERTEX=true
-MEDIC_VERTEX_PROJECT_ID=your-gcp-project-id
-MEDIC_VERTEX_LOCATION=us-central1
-```
-
 ### 3. Authenticate with HuggingFace
 
 ```bash
@@ -154,7 +140,7 @@ medic-amr-guard/
 ├── src/
 │   ├── agents.py           # Four agent implementations
 │   ├── graph.py            # LangGraph orchestrator + conditional routing
-│   ├── loader.py           # Model loading: local HuggingFace or Vertex AI
+│   ├── loader.py           # Model loading: local HuggingFace causal LMs
 │   ├── prompts.py          # System and user prompts for all agents
 │   ├── rag.py              # ChromaDB ingestion and retrieval helpers
 │   ├── state.py            # InfectionState TypedDict schema

notebooks/kaggle_medic_demo.ipynb

@@ -232,35 +232,7 @@
    "id": "a61f1fb1",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "# Write .env\n",
-    "env = f\"\"\"\n",
-    "MEDIC_ENV=kaggle\n",
-    "MEDIC_DEFAULT_BACKEND=local\n",
-    "MEDIC_USE_VERTEX=false\n",
-    "MEDIC_QUANTIZATION=4bit\n",
-    "\n",
-    "# Agent 1, 2, 4 — MedGemma 4B IT\n",
-    "MEDIC_LOCAL_MEDGEMMA_4B_MODEL={MEDGEMMA_4B}\n",
-    "\n",
-    "# Agent 3 — MedGemma 27B Text IT  (subbed with 4B for Kaggle T4)\n",
-    "# To use full 27B: set to google/medgemma-27b-text-it\n",
-    "MEDIC_LOCAL_MEDGEMMA_27B_MODEL={MEDGEMMA_4B}\n",
-    "\n",
-    "# Agent 4 safety — TxGemma 9B  (subbed with 2B for Kaggle T4)\n",
-    "# To use full 9B: set to google/txgemma-9b-predict\n",
-    "MEDIC_LOCAL_TXGEMMA_9B_MODEL={TXGEMMA_2B}\n",
-    "MEDIC_LOCAL_TXGEMMA_2B_MODEL={TXGEMMA_2B}\n",
-    "\n",
-    "MEDIC_EMBEDDING_MODEL=sentence-transformers/all-MiniLM-L6-v2\n",
-    "MEDIC_DATA_DIR=/kaggle/working/AMR-Guard/data\n",
-    "MEDIC_CHROMA_DB_DIR=/kaggle/working/AMR-Guard/data/chroma_db\n",
-    "\"\"\".strip()\n",
-    "\n",
-    "with open(\"/kaggle/working/AMR-Guard/.env\", \"w\") as f:\n",
-    "    f.write(env)\n",
-    "print(\".env written\")"
-   ]
+   "source": "# Write .env\nenv = f\"\"\"\nMEDIC_ENV=kaggle\nMEDIC_QUANTIZATION=4bit\n\n# Agent 1, 2, 4 — MedGemma 4B IT\nMEDIC_LOCAL_MEDGEMMA_4B_MODEL={MEDGEMMA_4B}\n\n# Agent 3 — MedGemma 27B Text IT  (subbed with 4B for Kaggle T4)\n# To use full 27B: set to google/medgemma-27b-text-it\nMEDIC_LOCAL_MEDGEMMA_27B_MODEL={MEDGEMMA_4B}\n\n# Agent 4 safety — TxGemma 9B  (subbed with 2B for Kaggle T4)\n# To use full 9B: set to google/txgemma-9b-predict\nMEDIC_LOCAL_TXGEMMA_9B_MODEL={TXGEMMA_2B}\nMEDIC_LOCAL_TXGEMMA_2B_MODEL={TXGEMMA_2B}\n\nMEDIC_EMBEDDING_MODEL=sentence-transformers/all-MiniLM-L6-v2\nMEDIC_DATA_DIR=/kaggle/working/AMR-Guard/data\nMEDIC_CHROMA_DB_DIR=/kaggle/working/AMR-Guard/data/chroma_db\n\"\"\".strip()\n\nwith open(\"/kaggle/working/AMR-Guard/.env\", \"w\") as f:\n    f.write(env)\nprint(\".env written\")"
   },
   {
    "cell_type": "code",
@@ -683,4 +655,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}

src/config.py

@@ -18,7 +18,6 @@ class Settings(BaseModel):
     All configuration for AMR-Guard, read from environment variables.
 
     Supports three deployment targets via MEDIC_ENV: local, kaggle, production.
-    Backend selection (vertex or local) is controlled by MEDIC_DEFAULT_BACKEND.
     """
 
     environment: Literal["local", "kaggle", "production"] = Field(
@@ -34,10 +33,7 @@ class Settings(BaseModel):
         default_factory=lambda: Path(os.getenv("MEDIC_CHROMA_DB_DIR", "data/chroma_db"))
     )
 
-    default_backend: Literal["vertex", "local"] = Field(
-        default_factory=lambda: os.getenv("MEDIC_DEFAULT_BACKEND", "local")  # type: ignore[arg-type]
-    )
-    # 4-bit quantization via bitsandbytes (local backend only)
+    # 4-bit quantization via bitsandbytes
     quantization: Literal["none", "4bit"] = Field(
         default_factory=lambda: os.getenv("MEDIC_QUANTIZATION", "4bit")  # type: ignore[arg-type]
     )
@@ -45,47 +41,17 @@ class Settings(BaseModel):
         default_factory=lambda: os.getenv("MEDIC_EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
     )
 
-    # Vertex AI settings
-    use_vertex: bool = Field(
-        default_factory=lambda: os.getenv("MEDIC_USE_VERTEX", "true").lower() in {"1", "true", "yes"}
-    )
-    vertex_project_id: Optional[str] = Field(
-        default_factory=lambda: os.getenv("MEDIC_VERTEX_PROJECT_ID")
-    )
-    vertex_location: str = Field(
-        default_factory=lambda: os.getenv("MEDIC_VERTEX_LOCATION", "us-central1")
-    )
-    vertex_medgemma_4b_model: str = Field(
-        default_factory=lambda: os.getenv("MEDIC_VERTEX_MEDGEMMA_4B_MODEL", "med-gemma-4b-it")
-    )
-    vertex_medgemma_27b_model: str = Field(
-        default_factory=lambda: os.getenv("MEDIC_VERTEX_MEDGEMMA_27B_MODEL", "med-gemma-27b-text-it")
-    )
-    vertex_txgemma_9b_model: str = Field(
-        default_factory=lambda: os.getenv("MEDIC_VERTEX_TXGEMMA_9B_MODEL", "tx-gemma-9b")
-    )
-    vertex_txgemma_2b_model: str = Field(
-        default_factory=lambda: os.getenv("MEDIC_VERTEX_TXGEMMA_2B_MODEL", "tx-gemma-2b")
-    )
-    google_application_credentials: Optional[Path] = Field(
-        default_factory=lambda: (
-            Path(os.environ["GOOGLE_APPLICATION_CREDENTIALS"])
-            if "GOOGLE_APPLICATION_CREDENTIALS" in os.environ
-            else None
-        )
-    )
-
-    # Local HuggingFace model paths (used when MEDIC_DEFAULT_BACKEND=local)
-    local_medgemma_4b_model: Optional[str] = Field(
+    # Local HuggingFace model paths
+    medgemma_4b_model: Optional[str] = Field(
         default_factory=lambda: os.getenv("MEDIC_LOCAL_MEDGEMMA_4B_MODEL")
     )
-    local_medgemma_27b_model: Optional[str] = Field(
+    medgemma_27b_model: Optional[str] = Field(
         default_factory=lambda: os.getenv("MEDIC_LOCAL_MEDGEMMA_27B_MODEL")
     )
-    local_txgemma_9b_model: Optional[str] = Field(
+    txgemma_9b_model: Optional[str] = Field(
         default_factory=lambda: os.getenv("MEDIC_LOCAL_TXGEMMA_9B_MODEL")
     )
-    local_txgemma_2b_model: Optional[str] = Field(
+    txgemma_2b_model: Optional[str] = Field(
         default_factory=lambda: os.getenv("MEDIC_LOCAL_TXGEMMA_2B_MODEL")
     )
 
@@ -94,4 +60,3 @@ class Settings(BaseModel):
 def get_settings() -> Settings:
     """Return the cached Settings singleton. Import this instead of instantiating Settings directly."""
     return Settings()
-

src/loader.py

@@ -7,58 +7,9 @@ from .config import get_settings
 
 logger = logging.getLogger(__name__)
 
-TextBackend = Literal["vertex", "local"]
 TextModelName = Literal["medgemma_4b", "medgemma_27b", "txgemma_9b", "txgemma_2b"]
 
 
-def _resolve_backend(requested: Optional[TextBackend]) -> TextBackend:
-    settings = get_settings()
-    backend = requested or settings.default_backend  # type: ignore[assignment]
-    if backend == "vertex" and not settings.use_vertex:
-        logger.info("Vertex disabled in settings, falling back to local backend.")
-        return "local"
-    return backend
-
-
-@lru_cache(maxsize=8)
-def _get_vertex_chat_model(model_name: TextModelName):
-    """Load a Vertex AI chat model and return a callable that takes a prompt string."""
-    try:
-        from langchain_google_vertexai import ChatVertexAI
-    except Exception as exc:
-        raise RuntimeError(
-            "langchain-google-vertexai is not available; "
-            "install it or switch MEDIC_DEFAULT_BACKEND=local."
-        ) from exc
-
-    settings = get_settings()
-    if settings.vertex_project_id is None:
-        raise RuntimeError(
-            "MEDIC_VERTEX_PROJECT_ID is not set. "
-            "Set it in your environment or .env to use Vertex AI."
-        )
-
-    model_id_map: Dict[TextModelName, str] = {
-        "medgemma_4b": settings.vertex_medgemma_4b_model,
-        "medgemma_27b": settings.vertex_medgemma_27b_model,
-        "txgemma_9b": settings.vertex_txgemma_9b_model,
-        "txgemma_2b": settings.vertex_txgemma_2b_model,
-    }
-
-    llm = ChatVertexAI(
-        model=model_id_map[model_name],
-        project=settings.vertex_project_id,
-        location=settings.vertex_location,
-        temperature=0.2,
-    )
-
-    def _call(prompt: str, **kwargs: Any) -> str:
-        result = llm.invoke(prompt, **kwargs)
-        return str(getattr(result, "content", result))
-
-    return _call
-
-
 @lru_cache(maxsize=8)
 def _get_local_causal_lm(model_name: TextModelName):
     """Load a local HuggingFace causal LM and return a generation callable."""
@@ -67,17 +18,17 @@ def _get_local_causal_lm(model_name: TextModelName):
 
     settings = get_settings()
     model_path_map: Dict[TextModelName, Optional[str]] = {
-        "medgemma_4b": settings.local_medgemma_4b_model,
-        "medgemma_27b": settings.local_medgemma_27b_model,
-        "txgemma_9b": settings.local_txgemma_9b_model,
-        "txgemma_2b": settings.local_txgemma_2b_model,
+        "medgemma_4b": settings.medgemma_4b_model,
+        "medgemma_27b": settings.medgemma_27b_model,
+        "txgemma_9b": settings.txgemma_9b_model,
+        "txgemma_2b": settings.txgemma_2b_model,
     }
 
     model_path = model_path_map[model_name]
     if not model_path:
         raise RuntimeError(
             f"No local model path configured for {model_name}. "
-            "Set MEDIC_LOCAL_*_MODEL or use the Vertex backend."
+            f"Set MEDIC_LOCAL_*_MODEL in your environment or .env."
         )
 
     load_kwargs: Dict[str, Any] = {"device_map": "auto"}
@@ -108,22 +59,18 @@ def _get_local_causal_lm(model_name: TextModelName):
 @lru_cache(maxsize=32)
 def get_text_model(
     model_name: TextModelName = "medgemma_4b",
-    backend: Optional[TextBackend] = None,
 ) -> Callable[..., str]:
-    """Return a cached callable for the requested model and backend."""
-    resolved = _resolve_backend(backend)
-    return _get_vertex_chat_model(model_name) if resolved == "vertex" else _get_local_causal_lm(model_name)
+    """Return a cached callable for the requested model."""
+    return _get_local_causal_lm(model_name)
 
 
 def run_inference(
     prompt: str,
     model_name: TextModelName = "medgemma_4b",
-    backend: Optional[TextBackend] = None,
     max_new_tokens: int = 512,
     temperature: float = 0.2,
     **kwargs: Any,
 ) -> str:
     """Run inference with the specified model. This is the primary entry point for agents."""
-    model = get_text_model(model_name=model_name, backend=backend)
+    model = get_text_model(model_name=model_name)
     return model(prompt, max_new_tokens=max_new_tokens, temperature=temperature, **kwargs)
-