feat: add grok search for general knowledge

Author: enitrat

.env.example

@@ -24,6 +24,7 @@ ANTHROPIC_API_KEY=""
 GEMINI_API_KEY=""
 DEEPSEEK_API_KEY=""
 GROQ_API_KEY=""
+XAI_API_KEY=""
 
 # Version Configuration
 STARKNET_FOUNDRY_VERSION="0.47.0"

python/pyproject.toml

@@ -52,6 +52,7 @@ dependencies = [
   "toml>=0.10.2",
   "tqdm>=4.66.0",
   "typer>=0.19.2",
+  "xai_sdk>=1.3.1",
 ]
 
 [project.optional-dependencies]

python/src/cairo_coder/core/rag_pipeline.py

@@ -26,6 +26,7 @@
 )
 from cairo_coder.dspy.document_retriever import DocumentRetrieverProgram
 from cairo_coder.dspy.generation_program import GenerationProgram, McpGenerationProgram
+from cairo_coder.dspy.grok_search import GrokSearchProgram
 from cairo_coder.dspy.query_processor import QueryProcessorProgram
 from cairo_coder.dspy.retrieval_judge import RetrievalJudge
 
@@ -73,6 +74,8 @@ def __init__(self, config: RagPipelineConfig):
         self.generation_program = config.generation_program
         self.mcp_generation_program = config.mcp_generation_program
         self.retrieval_judge = RetrievalJudge()
+        self.grok_search = GrokSearchProgram()
+        self._grok_citations: list[str] = []
 
         # Pipeline state
         self._current_processed_query: ProcessedQuery | None = None
@@ -96,6 +99,22 @@ async def _aprocess_query_and_retrieve_docs(
             processed_query=processed_query, sources=retrieval_sources
         )
 
+        # Optional Grok web/X augmentation: activate when STARKNET_BLOG is among sources.
+        try:
+            if DocumentSource.STARKNET_BLOG in retrieval_sources:
+                grok_docs = await self.grok_search.aforward(processed_query)
+                self._grok_citations = list(self.grok_search.last_citations)
+                if grok_docs:
+                    documents.extend(grok_docs)
+                grok_summary_doc = next((d for d in grok_docs if d.metadata.get("name") == "grok-answer"), None)
+            else:
+                self._grok_citations = []
+                grok_summary_doc = None
+        except Exception as e:
+            logger.warning("Grok augmentation failed; continuing without it", error=str(e), exc_info=True)
+            grok_summary_doc = None
+            self._grok_citations = []
+
         try:
             with dspy.context(
                 lm=dspy.LM("gemini/gemini-flash-lite-latest", max_tokens=10000, temperature=0.5),
@@ -110,6 +129,16 @@ async def _aprocess_query_and_retrieve_docs(
             )
             # documents already contains all retrieved docs, no action needed
 
+        # Ensure Grok summary is present and first in order (for generation context)
+        try:
+            if grok_summary_doc is not None:
+                if grok_summary_doc in documents:
+                    documents = [grok_summary_doc] + [d for d in documents if d is not grok_summary_doc]
+                else:
+                    documents = [grok_summary_doc] + documents
+        except Exception:
+            pass
+
         self._current_documents = documents
 
         return processed_query, documents
@@ -290,14 +319,34 @@ def _format_sources(self, documents: list[Document]) -> list[dict[str, Any]]:
             List of dicts: [{"title": str, "url": str}, ...]
         """
         sources: list[dict[str, str]] = []
+
+        # Helper to extract domain title
+        def title_from_url(url: str) -> str:
+            try:
+                import urllib.parse as _up
+
+                host = _up.urlparse(url).netloc
+                return host or url
+            except Exception:
+                return url
+
+        # 1) Vector store and other docs (skip Grok summary virtual doc)
         for doc in documents:
+            if doc.metadata.get("name") == "grok-answer" or doc.metadata.get("is_virtual"):
+                continue
             if doc.source_link is None:
                 logger.warning(f"Document {doc.title} has no source link")
-                to_append = ({"metadata": {"title": doc.title, "url": ""}})
+                to_append = {"metadata": {"title": doc.title, "url": ""}}
             else:
-                to_append = ({"metadata": {"title": doc.title, "url": doc.source_link}})
+                to_append = {"metadata": {"title": doc.title, "url": doc.source_link}}
             sources.append(to_append)
 
+        # 2) Append Grok citations (raw URLs)
+        for url in self._grok_citations:
+            if not url:
+                continue
+            sources.append({"metadata": {"title": title_from_url(url), "url": url}})
+
         return sources
 
     def _prepare_context(self, documents: list[Document]) -> str:
@@ -325,11 +374,12 @@ def _prepare_context(self, documents: list[Document]) -> str:
         for i, doc in enumerate(documents, 1):
             source_name = doc.metadata.get("source_display", "Unknown Source")
             title = doc.metadata.get("title", f"Document {i}")
-            url = doc.metadata.get("url", "#")
+            url = doc.metadata.get("url")
 
             context_parts.append(f"## {i}. {title}")
             context_parts.append(f"Source: {source_name}")
-            context_parts.append(f"URL: {url}")
+            if url:
+                context_parts.append(f"URL: {url}")
             context_parts.append("")
             context_parts.append(doc.page_content)
             context_parts.append("")

python/src/cairo_coder/dspy/__init__.py

@@ -15,6 +15,7 @@
     create_generation_program,
     create_mcp_generation_program,
 )
+from .grok_search import GrokSearchProgram
 from .query_processor import QueryProcessorProgram, create_query_processor
 from .retrieval_judge import RetrievalJudge
 from .suggestion_program import SuggestionGeneration
@@ -29,4 +30,5 @@
     "create_mcp_generation_program",
     "RetrievalJudge",
     "SuggestionGeneration",
+    "GrokSearchProgram",
 ]

python/src/cairo_coder/dspy/grok_search.py

@@ -0,0 +1,121 @@
+"""
+Grok Web/X Search module for Cairo Coder.
+
+Uses the xAI SDK (agentic server‑side tools: web_search, x_search) to fetch
+fresh context and a synthesized answer. The output is provided as a single
+virtual Document for the generator and a list of citation URLs that the
+pipeline emits via SOURCES.
+
+Behavior:
+- Activated upstream when DocumentSource.STARKNET_BLOG is in the requested sources.
+- Returns one primary virtual Document containing the Grok-composed answer
+  plus an inline source list inside the content.
+- Does not create per-citation documents; citations are emitted via SOURCES.
+
+Environment:
+- Set XAI_API_KEY with a valid xAI API key.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import os
+from urllib.parse import urlparse
+
+import dspy
+import structlog
+from xai_sdk import AsyncClient as XaiClient
+from xai_sdk.chat import Response, user
+from xai_sdk.tools import web_search, x_search
+
+from cairo_coder.core.types import Document, DocumentSource, ProcessedQuery
+
+logger = structlog.get_logger(__name__)
+
+
+DEFAULT_GROK_MODEL = "grok-4-fast"
+
+
+def _sha1(text: str) -> str:
+    return hashlib.sha1(text.encode("utf-8")).hexdigest()
+
+
+def _mk_unique_id(prefix: str, content: str, idx: int = 0) -> str:
+    return f"{prefix}-{_sha1(content)[:10]}-{idx}"
+
+
+
+class GrokSearchProgram(dspy.Module):
+    """
+    DSPy module that queries xAI's Grok Responses API with web and X search tools.
+
+    aforward returns a list[Document] suitable for inclusion in the RAG pipeline.
+    """
+
+    def __init__(
+        self,
+    ) -> None:
+        super().__init__()
+        api_key = os.getenv("XAI_API_KEY")
+        if not api_key:
+            raise RuntimeError("XAI_API_KEY must be set for GrokSearchProgram")
+        self.client = XaiClient(api_key=api_key)
+        self.last_citations: list[str] = []
+
+    @staticmethod
+    def _domain_from_url(url: str) -> str:
+        try:
+            return urlparse(url).netloc or url
+        except Exception:
+            return url
+
+    async def aforward(self, processed_query: ProcessedQuery) -> list[Document]:
+        formatted_query = f"""Answer the following query: {processed_query.original}. \
+            For more context, here are some semantic terms associated with the question: \
+            {', '.join(processed_query.search_queries)}.
+        """
+        chat = self.client.chat.create(
+            model=DEFAULT_GROK_MODEL,
+            tools=[web_search(), x_search()],
+        )
+        logger.info(f"Formatted query: {formatted_query}")
+        chat.append(user(formatted_query))
+        response: Response = await chat.sample()
+        answer: str = response.content
+        citations_urls: list[str] = response.citations
+        self.last_citations = list(citations_urls or [])
+        logger.info(f"Answer: {answer}")
+        logger.info(f"Citations URLs: {citations_urls}")
+
+        # Assemble a compact source list at the end
+        cite_lines = []
+        for i, url in enumerate(citations_urls):
+            title = self._domain_from_url(url) or f"Source {i}"
+            cite_lines.append(f"[{i}] {title}: {url}")
+        if cite_lines:
+            answer_with_sources = f"{answer}\n\nSources:\n" + "\n".join(cite_lines)
+        else:
+            answer_with_sources = answer
+
+        documents: list[Document] = []
+        unique_id = _mk_unique_id("grok-answer", answer)
+        documents.append(
+            Document(
+                page_content=answer_with_sources,
+                metadata={
+                    "name": "grok-answer",
+                    "title": "Grok Web/X Summary",
+                    "uniqueId": unique_id,
+                    "contentHash": _sha1(answer_with_sources),
+                    "chunkNumber": 0,
+                    # Treat as Starknet blog related to gate activation
+                    "source": DocumentSource.STARKNET_BLOG,
+                    "source_display": "Grok Web/X",
+                    "sourceLink": "",
+                    "url": "",
+                    "is_virtual": True,
+                },
+            )
+        )
+
+        return documents

python/tests/conftest.py

@@ -14,6 +14,7 @@
 
 from cairo_coder.core.agent_factory import AgentFactory
 from cairo_coder.core.config import VectorStoreConfig
+from cairo_coder.core.rag_pipeline import RagPipeline, RagPipelineConfig
 from cairo_coder.core.types import (
     Document,
     DocumentSource,
@@ -29,7 +30,7 @@
 from cairo_coder.server.app import CairoCoderServer, get_agent_factory
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(scope="function")
 def mock_returned_documents(sample_documents):
     """DSPy Examples derived from sample_documents for DRY content."""
     return [dspy.Example(content=doc.page_content, metadata=doc.metadata) for doc in sample_documents]
@@ -251,7 +252,7 @@ def sample_processed_query():
     )
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(scope="function")
 def sample_documents():
     """
     Create a collection of sample documents for testing.
@@ -334,6 +335,9 @@ def clean_config_env_vars(monkeypatch):
         original_values[var] = os.environ.get(var)
         monkeypatch.delenv(var, raising=False)
 
+    # Ensure xAI SDK clients can initialize in tests (no real network calls occur).
+    monkeypatch.setenv("XAI_API_KEY", "test")
+
     yield
 
     # Restore original values after test
@@ -450,3 +454,41 @@ async def async_filter_docs(query: str, documents: list[Document]) -> list[Docum
     judge.get_lm_usage = Mock(return_value={})
 
     return judge
+
+@pytest.fixture
+def pipeline_config(
+    mock_vector_store_config,
+    mock_query_processor,
+    mock_document_retriever,
+    mock_generation_program,
+    mock_mcp_generation_program,
+):
+    """Create a pipeline configuration."""
+    return RagPipelineConfig(
+        name="test_pipeline",
+        vector_store_config=mock_vector_store_config,
+        query_processor=mock_query_processor,
+        document_retriever=mock_document_retriever,
+        generation_program=mock_generation_program,
+        mcp_generation_program=mock_mcp_generation_program,
+        sources=list(DocumentSource),
+        max_source_count=10,
+        similarity_threshold=0.4,
+    )
+
+
+
+@pytest.fixture(scope="function")
+def pipeline(pipeline_config):
+    """Create a RagPipeline instance."""
+    with patch("cairo_coder.core.rag_pipeline.RetrievalJudge") as mock_judge_class:
+        mock_judge = Mock()
+        mock_judge.get_lm_usage.return_value = {}
+        mock_judge.aforward = AsyncMock(side_effect=lambda query, documents: documents)
+        mock_judge_class.return_value = mock_judge
+        return RagPipeline(pipeline_config)
+
+@pytest.fixture(scope="function")
+def rag_pipeline(pipeline_config):
+    """Alias fixture for pipeline to maintain backward compatibility."""
+    return RagPipeline(pipeline_config)

python/tests/unit/test_document_retriever.py

@@ -29,7 +29,7 @@ def retriever(
             similarity_threshold=0.4,
         )
 
-    @pytest.fixture(scope="session")
+    @pytest.fixture(scope="function")
     def mock_dspy_examples(self, sample_documents: list[Document]) -> list[dspy.Example]:
         """Create mock DSPy Example objects from sample documents."""
         examples = []