Sync from GitHub via hub-sync

README.md

@@ -16,6 +16,6 @@ Behavior:
 - Clones a public GitHub repo
 - Chunks it with tree-sitter
 - Builds retrieval state with a Qdrant adapter
-- Answers questions with Groq-hosted Llama or Vertex AI Gemini depending on environment configuration
+- Answers questions with Groq-hosted Llama, AWS Bedrock, or Vertex AI Gemini depending on environment configuration
 - Deletes the cloned repo after indexing
 - Keeps only lightweight repo metadata in SQLite

requirements.txt

@@ -5,6 +5,7 @@ pydantic==2.6.1
 python-dotenv==1.0.1
 
 openai==1.109.1
+boto3==1.40.58
 google-genai==1.12.1
 httpx==0.28.1
 numpy==1.26.4

src/embeddings.py

@@ -1,3 +1,4 @@
+import json
 import os
 import time
 from typing import Callable, List, Optional
@@ -71,6 +72,21 @@ class EmbeddingGenerator:
                     str(self.vertex_output_dimensionality or 3072),
                 )
             )
+        elif self.provider == "bedrock":
+            print(
+                f"[embeddings] Initializing AWS Bedrock embeddings with model={self.model_name}",
+                flush=True,
+            )
+            try:
+                import boto3
+            except ImportError as exc:
+                raise RuntimeError(
+                    "AWS Bedrock embedding support requires the `boto3` package."
+                ) from exc
+
+            region = os.getenv("AWS_REGION", os.getenv("AWS_DEFAULT_REGION", "us-east-1"))
+            self.client = boto3.client("bedrock-runtime", region_name=region)
+            self.embedding_dim = int(os.getenv("BEDROCK_EMBEDDING_DIM", "1024"))
         else:
             model_device = self.device or "cpu"
             print(
@@ -98,6 +114,8 @@ class EmbeddingGenerator:
                 [text],
                 task_type=self.vertex_task_type_query,
             )[0]
+        if self.provider == "bedrock":
+            return self._embed_with_bedrock(text)
         query_text = f"{self.query_prefix}: {text}" if self.query_prefix else text
         return self._encode_with_backoff([query_text], prompt_name=self.query_prompt_name)[0]
 
@@ -125,6 +143,12 @@ class EmbeddingGenerator:
                 batch_size=batch_size,
                 progress_callback=progress_callback,
             )
+        if self.provider == "bedrock":
+            return self._embed_batch_with_bedrock(
+                texts=texts,
+                batch_size=batch_size,
+                progress_callback=progress_callback,
+            )
 
         effective_batch_size = max(1, batch_size or self.batch_size)
         all_embeddings = []
@@ -222,6 +246,58 @@ class EmbeddingGenerator:
 
         return np.array(values, dtype="float32")
 
+    def _embed_batch_with_bedrock(
+        self,
+        texts: List[str],
+        batch_size: int = None,
+        progress_callback: Optional[Callable[[int, int], None]] = None,
+    ) -> np.ndarray:
+        effective_batch_size = max(1, batch_size or self.batch_size)
+        all_embeddings = []
+        total = len(texts)
+
+        for start in range(0, total, effective_batch_size):
+            batch = texts[start : start + effective_batch_size]
+            batch_number = (start // effective_batch_size) + 1
+            total_batches = (total + effective_batch_size - 1) // effective_batch_size
+            print(
+                f"[embeddings] Bedrock batch {batch_number}/{total_batches} "
+                f"items={len(batch)} progress={start}/{total}",
+                flush=True,
+            )
+            started_at = time.perf_counter()
+            batch_embeddings = [self._embed_with_bedrock(text) for text in batch]
+            all_embeddings.append(np.vstack(batch_embeddings))
+            elapsed = time.perf_counter() - started_at
+            print(
+                f"[embeddings] Finished Bedrock batch {batch_number}/{total_batches} "
+                f"elapsed={elapsed:.2f}s progress={min(start + len(batch), total)}/{total}",
+                flush=True,
+            )
+            if progress_callback:
+                progress_callback(min(start + len(batch), total), total)
+
+        return np.vstack(all_embeddings).astype("float32")
+
+    def _embed_with_bedrock(self, text: str) -> np.ndarray:
+        payload = {"inputText": text, "normalize": True}
+        if self.embedding_dim in {256, 512, 1024}:
+            payload["dimensions"] = self.embedding_dim
+
+        response = self.client.invoke_model(
+            modelId=self.model_name,
+            body=json.dumps(payload),
+            accept="application/json",
+            contentType="application/json",
+        )
+        body = json.loads(response["body"].read())
+        values = body.get("embedding")
+        if values is None:
+            values = (body.get("embeddingsByType") or {}).get("float")
+        if not values:
+            raise RuntimeError("AWS Bedrock embeddings returned an empty response.")
+        return np.array(values, dtype="float32")
+
     def _encode_with_backoff(
         self,
         texts: List[str],
@@ -266,12 +342,14 @@ class EmbeddingGenerator:
             return configured_provider
         if self._is_hf_space() or self._is_test_context():
             return "local"
-        return "vertex_ai"
+        return "bedrock"
 
     def _resolve_model_name(self) -> str:
         explicit_model = os.getenv("EMBEDDING_MODEL")
         if explicit_model:
             return explicit_model
+        if self.provider == "bedrock":
+            return os.getenv("BEDROCK_EMBEDDING_MODEL", "amazon.titan-embed-text-v2:0")
         if self.provider == "vertex_ai":
             return os.getenv("VERTEX_EMBEDDING_MODEL", "gemini-embedding-001")
         if self._is_hf_space() or self._is_test_context():

src/rag_system.py

@@ -45,7 +45,7 @@ class CodebaseRAGSystem:
             )
         )
         self.app_env = os.getenv("APP_ENV", os.getenv("ENVIRONMENT", "local")).lower()
-        self.llm_provider = os.getenv("LLM_PROVIDER", "vertex_ai").lower()
+        self.llm_provider = os.getenv("LLM_PROVIDER", "bedrock").lower()
         self.llm_client = None
         self.llm_model = ""
         self._configure_llm()
@@ -534,6 +534,21 @@ Do not leave the answer unfinished.
             self.llm_model = os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile")
             return
 
+        if self.llm_provider == "bedrock":
+            try:
+                import boto3
+            except ImportError as exc:
+                raise RuntimeError(
+                    "AWS Bedrock LLM support requires the `boto3` package."
+                ) from exc
+
+            region = os.getenv("AWS_REGION", os.getenv("AWS_DEFAULT_REGION", "us-east-1"))
+            self.llm_client = boto3.client("bedrock-runtime", region_name=region)
+            self.llm_model = os.getenv(
+                "BEDROCK_LLM_MODEL", "us.meta.llama3-3-70b-instruct-v1:0"
+            )
+            return
+
         if self.llm_provider == "vertex_ai":
             try:
                 from google import genai
@@ -575,6 +590,31 @@ Do not leave the answer unfinished.
             finish_reason = getattr(response.choices[0], "finish_reason", "") or ""
             return self._normalize_markdown_answer(content), str(finish_reason)
 
+        if self.llm_provider == "bedrock":
+            response = self.llm_client.converse(
+                modelId=self.llm_model,
+                system=[{"text": system_prompt.strip()}],
+                messages=[
+                    {
+                        "role": "user",
+                        "content": [{"text": user_prompt.strip()}],
+                    }
+                ],
+                inferenceConfig={
+                    "temperature": 0.1,
+                    "maxTokens": 2200,
+                },
+            )
+            output_message = (response.get("output") or {}).get("message") or {}
+            content_blocks = output_message.get("content") or []
+            text = "".join(
+                block.get("text", "") for block in content_blocks if isinstance(block, dict)
+            )
+            if not text.strip():
+                raise RuntimeError("AWS Bedrock returned an empty response.")
+            stop_reason = response.get("stopReason", "") or ""
+            return self._normalize_markdown_answer(text), str(stop_reason)
+
         response = self.llm_client.models.generate_content(
             model=self.llm_model,
             contents=f"{system_prompt.strip()}\n\n{user_prompt.strip()}",