From 798f5c34a716c14884c04e5d3743ecfcff429f4f Mon Sep 17 00:00:00 2001 From: Giuseppe Zileni Date: Wed, 10 Jun 2026 14:29:21 +0200 Subject: [PATCH 1/2] feat(api): pluggable LLM provider (Ollama/Anthropic) + /ingest/json MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extraction-time LLM is now decoupled from the legacy Ollama-only path through a small ``LLMProvider`` ABC under ``pipeline/llm/``. The active provider is selected by ``KG_LLM_PROVIDER`` (default ``ollama``); set it to ``anthropic`` to route extraction to Claude Haiku via the official ``anthropic`` SDK. Embeddings stay on Ollama unconditionally — the vector index is sized for ``nomic-embed-text``. Also adds ``POST /ingest/json`` for ingesting structured news articles without writing files on the client side. The endpoint serialises ``{title, body, url, source, published_at, namespace}`` to a temp .txt and feeds it through the existing pipeline (no router changes needed — .txt is already supported). Tests: 14 new tests (extractor refactor, both providers, factory, JSON ingest serialisation contract). Full suite 23/23 green. --- knowledge-graph-api/.env.example | 8 ++ knowledge-graph-api/api/main.py | 11 +- knowledge-graph-api/api/routes/ingest.py | 75 +++++++++++- knowledge-graph-api/api/schemas.py | 28 ++++- knowledge-graph-api/config/settings.py | 10 ++ knowledge-graph-api/pipeline/extractor.py | 36 ++---- knowledge-graph-api/pipeline/llm/__init__.py | 18 +++ .../pipeline/llm/anthropic_provider.py | 72 +++++++++++ knowledge-graph-api/pipeline/llm/base.py | 27 +++++ knowledge-graph-api/pipeline/llm/factory.py | 22 ++++ .../pipeline/llm/ollama_provider.py | 33 +++++ knowledge-graph-api/requirements.txt | 5 +- knowledge-graph-api/tests/test_extractor.py | 87 ++++++++++---- knowledge-graph-api/tests/test_ingest_json.py | 89 ++++++++++++++ .../tests/test_llm_providers.py | 113 ++++++++++++++++++ 15 files changed, 586 insertions(+), 48 deletions(-) create mode 100644 knowledge-graph-api/pipeline/llm/__init__.py create mode 100644 knowledge-graph-api/pipeline/llm/anthropic_provider.py create mode 100644 knowledge-graph-api/pipeline/llm/base.py create mode 100644 knowledge-graph-api/pipeline/llm/factory.py create mode 100644 knowledge-graph-api/pipeline/llm/ollama_provider.py create mode 100644 knowledge-graph-api/tests/test_ingest_json.py create mode 100644 knowledge-graph-api/tests/test_llm_providers.py diff --git a/knowledge-graph-api/.env.example b/knowledge-graph-api/.env.example index 134ae47..78f0784 100644 --- a/knowledge-graph-api/.env.example +++ b/knowledge-graph-api/.env.example @@ -9,11 +9,19 @@ REDIS_URL=redis://redis:6379 REDIS_INDEX_NAME=kg_vectors REDIS_VECTOR_DIM=768 +# LLM provider for extraction stage: "ollama" (default) or "anthropic". +# Embeddings always go to Ollama regardless of this setting. +KG_LLM_PROVIDER=ollama + # Ollama (Inference API locale) OLLAMA_BASE_URL=http://ollama:11434 OLLAMA_LLM_MODEL=llama3 OLLAMA_EMBEDDING_MODEL=nomic-embed-text +# Anthropic (used only when KG_LLM_PROVIDER=anthropic) +ANTHROPIC_API_KEY= +ANTHROPIC_EXTRACTION_MODEL=claude-haiku-4-5 + # Chunking CHUNK_SIZE=1024 CHUNK_OVERLAP=128 diff --git a/knowledge-graph-api/api/main.py b/knowledge-graph-api/api/main.py index a357cf3..1ad044d 100644 --- a/knowledge-graph-api/api/main.py +++ b/knowledge-graph-api/api/main.py @@ -72,7 +72,9 @@ async def health() -> HealthResponse: except Exception as exc: logger.warning("health_redis_fail", error=str(exc)) - # Ollama + # LLM provider — meaning depends on KG_LLM_PROVIDER. When "anthropic" + # we still ping Ollama (because embeddings always run there); the + # field name stays "ollama" for backward compat with existing clients. try: async with httpx.AsyncClient(timeout=5.0) as client: resp = await client.get(f"{settings.OLLAMA_BASE_URL}/api/tags") @@ -80,12 +82,19 @@ async def health() -> HealthResponse: except Exception as exc: logger.warning("health_ollama_fail", error=str(exc)) + # When anthropic is configured as extraction provider, also surface + # whether the API key is present (cheap check — no network call). + provider = settings.KG_LLM_PROVIDER.lower().strip() + if provider == "anthropic" and not settings.ANTHROPIC_API_KEY: + logger.warning("health_anthropic_no_key") + all_ok = neo4j_ok and redis_ok and ollama_ok return HealthResponse( status="healthy" if all_ok else "degraded", neo4j=neo4j_ok, redis=redis_ok, ollama=ollama_ok, + llm_provider=provider, ) diff --git a/knowledge-graph-api/api/routes/ingest.py b/knowledge-graph-api/api/routes/ingest.py index 1599bdf..e357682 100644 --- a/knowledge-graph-api/api/routes/ingest.py +++ b/knowledge-graph-api/api/routes/ingest.py @@ -8,7 +8,7 @@ from fastapi import APIRouter, File, Form, HTTPException, UploadFile -from api.schemas import IngestRequest +from api.schemas import IngestJsonRequest, IngestRequest from pipeline.ingest import IngestOptions, IngestResult, IngestionPipeline from utils.logger import logger @@ -129,3 +129,76 @@ async def upload_and_ingest( pass return result + + +@router.post("/ingest/json", response_model=IngestResult) +async def ingest_json(body: IngestJsonRequest) -> IngestResult: + """Ingest a structured JSON document (e.g. a news article). + + Serialises the payload as a plain-text file with title, body, and + a trailing metadata block, then runs it through the standard + ingestion pipeline. The ``namespace`` field maps to ``thread_id`` + (the pipeline's partition key). + + Args: + body: The structured document to ingest. + + Returns: + An ``IngestResult`` summary. + """ + safe_stem = "".join( + c if c.isalnum() or c in "-_ " else "_" for c in body.title + )[:60].strip() or "article" + + serialised = ( + f"{body.title}\n\n" + f"{body.body}\n\n" + "---\n" + f"Source: {body.source}\n" + f"URL: {body.url}\n" + f"Published: {body.published_at.isoformat()}\n" + ) + + tmp_path: str | None = None + try: + with tempfile.NamedTemporaryFile( + delete=False, + suffix=".txt", + prefix=safe_stem + "_", + mode="w", + encoding="utf-8", + ) as tmp: + tmp.write(serialised) + tmp_path = tmp.name + + logger.info( + "ingest_json_received", + title=body.title, + source=body.source, + namespace=body.namespace, + size=len(serialised), + tmp=tmp_path, + ) + + pipeline = IngestionPipeline() + try: + result = await pipeline.ingest( + file_path=tmp_path, + thread_id=body.namespace, + options=IngestOptions(skip_existing=body.skip_existing), + ) + except ValueError as exc: + logger.error("ingest_json_error", error=str(exc)) + raise HTTPException(status_code=400, detail=str(exc)) + except Exception as exc: + logger.error("ingest_json_error", error=str(exc)) + raise HTTPException(status_code=500, detail="Ingestion failed") + + finally: + if tmp_path: + try: + os.unlink(tmp_path) + except OSError: + pass + + return result diff --git a/knowledge-graph-api/api/schemas.py b/knowledge-graph-api/api/schemas.py index 4348c35..eb7e3e6 100644 --- a/knowledge-graph-api/api/schemas.py +++ b/knowledge-graph-api/api/schemas.py @@ -2,7 +2,9 @@ from __future__ import annotations -from pydantic import BaseModel +from datetime import datetime + +from pydantic import BaseModel, ConfigDict, HttpUrl class IngestRequest(BaseModel): @@ -13,6 +15,26 @@ class IngestRequest(BaseModel): skip_existing: bool = True +class IngestJsonRequest(BaseModel): + """Body for POST /ingest/json — structured news ingestion. + + Designed for ingesting financial news (or any short structured + document) without writing a file to disk on the client side. The + server serialises the payload to a .txt file and feeds it through + the standard ingestion pipeline. + """ + + model_config = ConfigDict(extra="forbid") + + title: str + body: str + url: HttpUrl + source: str + published_at: datetime + namespace: str + skip_existing: bool = True + + class QueryRequest(BaseModel): """Body for POST /query and POST /query/stream.""" @@ -28,4 +50,8 @@ class HealthResponse(BaseModel): status: str neo4j: bool redis: bool + # ``ollama`` stays for backward compatibility with existing clients. + # When KG_LLM_PROVIDER=anthropic this field reports the Anthropic + # API reachability instead, and ``llm_provider`` identifies which. ollama: bool + llm_provider: str = "ollama" diff --git a/knowledge-graph-api/config/settings.py b/knowledge-graph-api/config/settings.py index 1908eb5..0de7225 100644 --- a/knowledge-graph-api/config/settings.py +++ b/knowledge-graph-api/config/settings.py @@ -17,6 +17,12 @@ class Settings(BaseSettings): REDIS_INDEX_NAME: str = "kg_vectors" REDIS_VECTOR_DIM: int = 768 + # LLM provider selection for the extraction stage. + # "ollama" (default) keeps the legacy local-inference behaviour. + # "anthropic" routes extraction to Claude (Haiku by default). + # Embeddings always go to Ollama regardless of this setting. + KG_LLM_PROVIDER: str = "ollama" + # Ollama (Inference API locale) OLLAMA_BASE_URL: str = "http://localhost:11434" OLLAMA_LLM_MODEL: str = "llama3" @@ -27,6 +33,10 @@ class Settings(BaseSettings): OLLAMA_EXTRACTION_MODEL: str = "" OLLAMA_EMBEDDING_MODEL: str = "nomic-embed-text" + # Anthropic (used when KG_LLM_PROVIDER=anthropic) + ANTHROPIC_API_KEY: str = "" + ANTHROPIC_EXTRACTION_MODEL: str = "claude-haiku-4-5" + # App CHUNK_SIZE: int = 1024 CHUNK_OVERLAP: int = 128 diff --git a/knowledge-graph-api/pipeline/extractor.py b/knowledge-graph-api/pipeline/extractor.py index b0921b1..25e6c15 100644 --- a/knowledge-graph-api/pipeline/extractor.py +++ b/knowledge-graph-api/pipeline/extractor.py @@ -1,16 +1,22 @@ -"""Entity and relation extraction via Ollama LLM.""" +"""Entity and relation extraction via a pluggable LLM provider. + +The provider (Ollama or Anthropic Claude) is resolved at construction +time via ``pipeline.llm.get_llm_provider()`` driven by the +``KG_LLM_PROVIDER`` setting. The extractor itself is provider-agnostic: +it sends the system + user prompts and parses the JSON response. +""" from __future__ import annotations import json -import httpx from pydantic import BaseModel from tenacity import retry, stop_after_attempt, wait_exponential from config.settings import settings from models.graph_node import VALID_NODE_TYPES from models.relation import VALID_RELATION_TYPES +from pipeline.llm import LLMProvider, get_llm_provider from utils.logger import logger # ── Result models ──────────────────────────────────────────────────── @@ -119,12 +125,10 @@ class ExtractionResult(BaseModel): # ── Extractor class ────────────────────────────────────────────────── class EntityExtractor: - """Extracts entities and relations from text chunks using Ollama.""" + """Extracts entities and relations from text chunks via an LLM provider.""" - def __init__(self) -> None: - self.base_url = settings.OLLAMA_BASE_URL - # Use dedicated extraction model if configured, otherwise fall back to LLM model. - self.model = settings.OLLAMA_EXTRACTION_MODEL or settings.OLLAMA_LLM_MODEL + def __init__(self, provider: LLMProvider | None = None) -> None: + self._provider = provider or get_llm_provider() async def extract(self, chunk: str, context: str = "") -> ExtractionResult: """Extract entities and relations from a text chunk. @@ -163,19 +167,5 @@ async def extract(self, chunk: str, context: str = "") -> ExtractionResult: reraise=True, ) async def _call_llm(self, user_content: str) -> str: - """Call the Ollama chat endpoint and return the raw message content.""" - async with httpx.AsyncClient(timeout=120.0) as client: - response = await client.post( - f"{self.base_url}/api/chat", - json={ - "model": self.model, - "messages": [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": user_content}, - ], - "stream": False, - "format": "json", - }, - ) - response.raise_for_status() - return response.json()["message"]["content"] + """Delegate the chat call to the configured LLM provider.""" + return await self._provider.chat_json(SYSTEM_PROMPT, user_content) diff --git a/knowledge-graph-api/pipeline/llm/__init__.py b/knowledge-graph-api/pipeline/llm/__init__.py new file mode 100644 index 0000000..d21967b --- /dev/null +++ b/knowledge-graph-api/pipeline/llm/__init__.py @@ -0,0 +1,18 @@ +"""LLM provider abstraction for entity/relation extraction. + +The extraction stage of the pipeline calls an LLM to convert text chunks +into structured JSON. The provider is pluggable so the same pipeline can +run against a local Ollama model, a hosted Anthropic Claude model, or +future providers without touching the call sites. + +Embeddings are NOT covered by this abstraction — they remain wired to +``OLLAMA_EMBEDDING_MODEL`` via ``pipeline.embedder`` because the vector +index is sized to that model's output (768 dims for ``nomic-embed-text``). +""" + +from __future__ import annotations + +from pipeline.llm.base import LLMProvider +from pipeline.llm.factory import get_llm_provider + +__all__ = ["LLMProvider", "get_llm_provider"] diff --git a/knowledge-graph-api/pipeline/llm/anthropic_provider.py b/knowledge-graph-api/pipeline/llm/anthropic_provider.py new file mode 100644 index 0000000..04321b8 --- /dev/null +++ b/knowledge-graph-api/pipeline/llm/anthropic_provider.py @@ -0,0 +1,72 @@ +"""Anthropic Claude provider for extraction-time JSON output. + +Uses the ``anthropic`` Python SDK directly (no MAF dependency) to keep +this package lightweight. Targets Haiku by default for cost: extraction +is a high-volume, narrow-schema task — Haiku 4.5 is the right tool. + +The provider forces JSON output by appending a brief reinforcement to +the system prompt and stripping any code fences from the response. +""" + +from __future__ import annotations + +from anthropic import AsyncAnthropic + +from config.settings import settings +from pipeline.llm.base import LLMProvider + +_JSON_REINFORCEMENT = ( + "\n\nIMPORTANT: respond with raw JSON only. No prose, no code fences, " + "no preamble — just the JSON object starting with { and ending with }." +) + + +def _strip_code_fences(text: str) -> str: + """Remove ```json … ``` wrappers that Claude sometimes emits anyway.""" + + stripped = text.strip() + if stripped.startswith("```"): + # drop the opening fence (with or without language tag) + first_newline = stripped.find("\n") + if first_newline != -1: + stripped = stripped[first_newline + 1 :] + # drop the closing fence + if stripped.rstrip().endswith("```"): + stripped = stripped.rstrip()[:-3] + return stripped.strip() + + +class AnthropicProvider(LLMProvider): + """Calls Claude via the Anthropic Messages API and returns raw JSON text.""" + + def __init__( + self, + *, + api_key: str | None = None, + model: str | None = None, + max_tokens: int = 4096, + ) -> None: + resolved_key = api_key or settings.ANTHROPIC_API_KEY + if not resolved_key: + raise RuntimeError( + "ANTHROPIC_API_KEY is required when KG_LLM_PROVIDER=anthropic" + ) + self._client = AsyncAnthropic(api_key=resolved_key) + self._model = model or settings.ANTHROPIC_EXTRACTION_MODEL + self._max_tokens = max_tokens + + async def chat_json(self, system: str, user: str) -> str: + response = await self._client.messages.create( + model=self._model, + max_tokens=self._max_tokens, + system=system + _JSON_REINFORCEMENT, + messages=[{"role": "user", "content": user}], + ) + # The Anthropic SDK returns a list of content blocks; for plain + # text completions there is exactly one TextBlock. + parts: list[str] = [] + for block in response.content: + text = getattr(block, "text", None) + if isinstance(text, str): + parts.append(text) + return _strip_code_fences("".join(parts)) diff --git a/knowledge-graph-api/pipeline/llm/base.py b/knowledge-graph-api/pipeline/llm/base.py new file mode 100644 index 0000000..a31acf0 --- /dev/null +++ b/knowledge-graph-api/pipeline/llm/base.py @@ -0,0 +1,27 @@ +"""Abstract LLM provider interface for extraction-time chat calls.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class LLMProvider(ABC): + """Provider contract for the extraction LLM. + + Implementations must return the raw JSON string emitted by the model + so the caller can ``json.loads`` it. Any parse-error tolerance lives + in the caller, not the provider — providers should NOT wrap parsing. + """ + + @abstractmethod + async def chat_json(self, system: str, user: str) -> str: + """Send a chat completion and return the raw response text. + + Args: + system: System prompt (instructions + schema + examples). + user: User content (the text chunk to extract from). + + Returns: + The model's raw textual output — expected to be JSON, but the + provider does not validate that. + """ diff --git a/knowledge-graph-api/pipeline/llm/factory.py b/knowledge-graph-api/pipeline/llm/factory.py new file mode 100644 index 0000000..17c551b --- /dev/null +++ b/knowledge-graph-api/pipeline/llm/factory.py @@ -0,0 +1,22 @@ +"""Factory that resolves the configured ``LLMProvider``.""" + +from __future__ import annotations + +from config.settings import settings +from pipeline.llm.anthropic_provider import AnthropicProvider +from pipeline.llm.base import LLMProvider +from pipeline.llm.ollama_provider import OllamaProvider + + +def get_llm_provider() -> LLMProvider: + """Return the provider selected by ``KG_LLM_PROVIDER`` (default: ollama).""" + + choice = settings.KG_LLM_PROVIDER.lower().strip() + if choice == "anthropic": + return AnthropicProvider() + if choice == "ollama": + return OllamaProvider() + raise ValueError( + f"Unknown KG_LLM_PROVIDER={settings.KG_LLM_PROVIDER!r} " + "(supported: 'ollama', 'anthropic')" + ) diff --git a/knowledge-graph-api/pipeline/llm/ollama_provider.py b/knowledge-graph-api/pipeline/llm/ollama_provider.py new file mode 100644 index 0000000..19b6451 --- /dev/null +++ b/knowledge-graph-api/pipeline/llm/ollama_provider.py @@ -0,0 +1,33 @@ +"""Ollama-backed LLM provider (legacy default).""" + +from __future__ import annotations + +import httpx + +from config.settings import settings +from pipeline.llm.base import LLMProvider + + +class OllamaProvider(LLMProvider): + """Calls the Ollama chat endpoint with ``format:"json"`` for structured output.""" + + def __init__(self, *, base_url: str | None = None, model: str | None = None) -> None: + self._base_url = (base_url or settings.OLLAMA_BASE_URL).rstrip("/") + self._model = model or settings.OLLAMA_EXTRACTION_MODEL or settings.OLLAMA_LLM_MODEL + + async def chat_json(self, system: str, user: str) -> str: + async with httpx.AsyncClient(timeout=120.0) as client: + response = await client.post( + f"{self._base_url}/api/chat", + json={ + "model": self._model, + "messages": [ + {"role": "system", "content": system}, + {"role": "user", "content": user}, + ], + "stream": False, + "format": "json", + }, + ) + response.raise_for_status() + return response.json()["message"]["content"] diff --git a/knowledge-graph-api/requirements.txt b/knowledge-graph-api/requirements.txt index d1f0cda..0a09682 100644 --- a/knowledge-graph-api/requirements.txt +++ b/knowledge-graph-api/requirements.txt @@ -1,6 +1,9 @@ -# HTTP client per Ollama +# HTTP client (Ollama + general purpose) httpx>=0.27.0 +# Anthropic Claude SDK (used when KG_LLM_PROVIDER=anthropic) +anthropic>=0.40.0 + # Neo4j neo4j>=5.18.0 diff --git a/knowledge-graph-api/tests/test_extractor.py b/knowledge-graph-api/tests/test_extractor.py index efa2fb8..cdcea79 100644 --- a/knowledge-graph-api/tests/test_extractor.py +++ b/knowledge-graph-api/tests/test_extractor.py @@ -2,43 +2,88 @@ from __future__ import annotations -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock import pytest from pipeline.extractor import EntityExtractor +from pipeline.llm.base import LLMProvider + + +class _StubProvider(LLMProvider): + def __init__(self, response: str) -> None: + self._response = response + + async def chat_json(self, system: str, user: str) -> str: + return self._response @pytest.mark.asyncio -async def test_extract_returns_result(mock_ollama_client) -> None: - """Extractor should return an ExtractionResult (possibly empty).""" - # Override the chat response to return valid JSON - chat_resp = MagicMock() - chat_resp.status_code = 200 - chat_resp.json.return_value = { - "message": { - "content": '{"entities": [{"id": "redis", "name": "Redis", "type": "Technology", "description": "Data store.", "importance": 0.8, "confidence": 0.9}], "relations": []}' - } - } - chat_resp.raise_for_status = MagicMock() - mock_ollama_client.post = AsyncMock(return_value=chat_resp) +async def test_extract_returns_result() -> None: + """Extractor should parse a well-formed JSON response into entities.""" + provider = _StubProvider( + '{"entities": [{"id": "redis", "name": "Redis", "type": "Technology",' + ' "description": "Data store.", "importance": 0.8, "confidence": 0.9}],' + ' "relations": []}' + ) - extractor = EntityExtractor() + extractor = EntityExtractor(provider=provider) result = await extractor.extract("Redis is an in-memory data store.") assert len(result.entities) == 1 assert result.entities[0].name == "Redis" @pytest.mark.asyncio -async def test_extract_handles_bad_json(mock_ollama_client) -> None: +async def test_extract_handles_bad_json() -> None: """Extractor should gracefully handle unparseable LLM output.""" - bad_resp = MagicMock() - bad_resp.status_code = 200 - bad_resp.json.return_value = {"message": {"content": "not json at all"}} - bad_resp.raise_for_status = MagicMock() - mock_ollama_client.post = AsyncMock(return_value=bad_resp) + provider = _StubProvider("not json at all") - extractor = EntityExtractor() + extractor = EntityExtractor(provider=provider) result = await extractor.extract("Some text") assert len(result.entities) == 0 assert len(result.relations) == 0 + + +@pytest.mark.asyncio +async def test_extract_filters_low_confidence() -> None: + """Entities below the 0.6 confidence threshold are dropped.""" + provider = _StubProvider( + '{"entities": [' + '{"id": "high", "name": "High", "type": "Technology", "description": "x",' + ' "importance": 0.5, "confidence": 0.9},' + '{"id": "low", "name": "Low", "type": "Technology", "description": "x",' + ' "importance": 0.5, "confidence": 0.4}' + '], "relations": []}' + ) + + extractor = EntityExtractor(provider=provider) + result = await extractor.extract("anything") + assert {e.name for e in result.entities} == {"High"} + + +@pytest.mark.asyncio +async def test_extractor_default_provider_is_resolved_from_factory(monkeypatch) -> None: + """Without an explicit provider, the extractor pulls one from the factory.""" + sentinel = _StubProvider('{"entities": [], "relations": []}') + called = {} + + def fake_factory(): + called["yes"] = True + return sentinel + + monkeypatch.setattr("pipeline.extractor.get_llm_provider", fake_factory) + captured: dict[str, str] = {} + + original = sentinel.chat_json + + async def spy(system: str, user: str) -> str: + captured["user"] = user + return await original(system, user) + + sentinel.chat_json = AsyncMock(side_effect=spy) + + extractor = EntityExtractor() + await extractor.extract("hello world") + + assert called == {"yes": True} + assert captured["user"] == "hello world" diff --git a/knowledge-graph-api/tests/test_ingest_json.py b/knowledge-graph-api/tests/test_ingest_json.py new file mode 100644 index 0000000..46e924e --- /dev/null +++ b/knowledge-graph-api/tests/test_ingest_json.py @@ -0,0 +1,89 @@ +"""Tests for the POST /ingest/json route.""" + +from __future__ import annotations + +import os +from datetime import datetime +from unittest.mock import AsyncMock, patch + +import pytest + +from api.routes.ingest import ingest_json +from api.schemas import IngestJsonRequest +from pipeline.ingest import IngestResult + + +@pytest.mark.asyncio +async def test_ingest_json_serialises_payload_and_invokes_pipeline() -> None: + """The handler should write a .txt with title+body+metadata, call the pipeline, and clean up.""" + + captured: dict[str, str] = {} + + async def fake_ingest(file_path: str, thread_id: str, options) -> IngestResult: + captured["file_path"] = file_path + captured["thread_id"] = thread_id + captured["skip_existing"] = options.skip_existing + with open(file_path, encoding="utf-8") as f: + captured["content"] = f.read() + return IngestResult( + document_id="doc_test", + chunks_processed=2, + chunks_skipped=0, + entities_extracted=0, + relations_extracted=0, + nodes_created=0, + edges_created=0, + processing_time_ms=10.0, + errors=[], + ) + + with patch("api.routes.ingest.IngestionPipeline") as cls: + cls.return_value.ingest = AsyncMock(side_effect=fake_ingest) + + body = IngestJsonRequest( + title="Apple beats earnings", + body="Q4 revenue up 5% YoY thanks to Services.", + url="https://reuters.example/article/123", + source="Reuters", + published_at=datetime(2026, 1, 23, 14, 30, 0), + namespace="news-us-eu", + ) + result = await ingest_json(body) + + assert result.document_id == "doc_test" + assert result.chunks_processed == 2 + assert captured["thread_id"] == "news-us-eu" + assert captured["skip_existing"] is True + + content = captured["content"] + assert content.startswith("Apple beats earnings\n\n") + assert "Q4 revenue up 5% YoY" in content + assert "Source: Reuters" in content + assert "https://reuters.example/article/123" in content + assert "Published: 2026-01-23T14:30:00" in content + + # Tempfile must have been unlinked. + assert not os.path.exists(captured["file_path"]) + + +@pytest.mark.asyncio +async def test_ingest_json_propagates_pipeline_value_error() -> None: + """ValueError from the pipeline becomes a 400.""" + from fastapi import HTTPException + + with patch("api.routes.ingest.IngestionPipeline") as cls: + cls.return_value.ingest = AsyncMock(side_effect=ValueError("unsupported mime")) + + body = IngestJsonRequest( + title="t", + body="b", + url="https://example.com/x", + source="X", + published_at=datetime(2026, 1, 1), + namespace="ns", + ) + with pytest.raises(HTTPException) as exc: + await ingest_json(body) + + assert exc.value.status_code == 400 + assert "unsupported mime" in exc.value.detail diff --git a/knowledge-graph-api/tests/test_llm_providers.py b/knowledge-graph-api/tests/test_llm_providers.py new file mode 100644 index 0000000..9975672 --- /dev/null +++ b/knowledge-graph-api/tests/test_llm_providers.py @@ -0,0 +1,113 @@ +"""Tests for the LLM provider abstraction (Ollama + Anthropic + factory).""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from config.settings import settings +from pipeline.llm.anthropic_provider import AnthropicProvider, _strip_code_fences +from pipeline.llm.factory import get_llm_provider +from pipeline.llm.ollama_provider import OllamaProvider + + +@pytest.mark.asyncio +async def test_ollama_provider_calls_chat_endpoint(mock_ollama_client) -> None: + """OllamaProvider POSTs to /api/chat and returns the message content.""" + resp = MagicMock() + resp.status_code = 200 + resp.json.return_value = {"message": {"content": '{"entities": []}'}} + resp.raise_for_status = MagicMock() + mock_ollama_client.post = AsyncMock(return_value=resp) + + provider = OllamaProvider(base_url="http://o:11434", model="llama3") + result = await provider.chat_json("sys", "user") + + assert result == '{"entities": []}' + call = mock_ollama_client.post.await_args + assert call.args[0] == "http://o:11434/api/chat" + payload = call.kwargs["json"] + assert payload["model"] == "llama3" + assert payload["format"] == "json" + assert payload["messages"][0] == {"role": "system", "content": "sys"} + assert payload["messages"][1] == {"role": "user", "content": "user"} + + +@pytest.mark.asyncio +async def test_anthropic_provider_calls_messages_api() -> None: + """AnthropicProvider invokes messages.create and returns concatenated text.""" + fake_block = MagicMock() + fake_block.text = '{"entities": []}' + fake_response = MagicMock() + fake_response.content = [fake_block] + + fake_client = MagicMock() + fake_client.messages.create = AsyncMock(return_value=fake_response) + + with patch( + "pipeline.llm.anthropic_provider.AsyncAnthropic", + return_value=fake_client, + ): + provider = AnthropicProvider(api_key="sk-test", model="claude-haiku-4-5") + text = await provider.chat_json("sys", "user") + + assert text == '{"entities": []}' + kwargs = fake_client.messages.create.await_args.kwargs + assert kwargs["model"] == "claude-haiku-4-5" + assert kwargs["messages"] == [{"role": "user", "content": "user"}] + # Reinforcement appended to enforce JSON-only output. + assert kwargs["system"].startswith("sys") + assert "raw JSON only" in kwargs["system"] + + +@pytest.mark.asyncio +async def test_anthropic_provider_strips_code_fences() -> None: + """Models occasionally wrap JSON in ```json … ``` — we strip them.""" + fake_block = MagicMock() + fake_block.text = '```json\n{"entities": []}\n```' + fake_response = MagicMock() + fake_response.content = [fake_block] + + fake_client = MagicMock() + fake_client.messages.create = AsyncMock(return_value=fake_response) + + with patch( + "pipeline.llm.anthropic_provider.AsyncAnthropic", + return_value=fake_client, + ): + provider = AnthropicProvider(api_key="sk-test") + text = await provider.chat_json("sys", "user") + + assert text == '{"entities": []}' + + +def test_strip_code_fences_passthrough() -> None: + """No fences = no change.""" + assert _strip_code_fences('{"a": 1}') == '{"a": 1}' + + +def test_anthropic_provider_rejects_missing_key(monkeypatch) -> None: + """Construction must fail loudly when the key is missing.""" + monkeypatch.setattr(settings, "ANTHROPIC_API_KEY", "") + with pytest.raises(RuntimeError, match="ANTHROPIC_API_KEY"): + AnthropicProvider() + + +def test_factory_returns_ollama_by_default(monkeypatch) -> None: + monkeypatch.setattr(settings, "KG_LLM_PROVIDER", "ollama") + assert isinstance(get_llm_provider(), OllamaProvider) + + +def test_factory_returns_anthropic_when_configured(monkeypatch) -> None: + monkeypatch.setattr(settings, "KG_LLM_PROVIDER", "anthropic") + monkeypatch.setattr(settings, "ANTHROPIC_API_KEY", "sk-test") + with patch("pipeline.llm.anthropic_provider.AsyncAnthropic"): + provider = get_llm_provider() + assert isinstance(provider, AnthropicProvider) + + +def test_factory_rejects_unknown_provider(monkeypatch) -> None: + monkeypatch.setattr(settings, "KG_LLM_PROVIDER", "wrong") + with pytest.raises(ValueError, match="Unknown KG_LLM_PROVIDER"): + get_llm_provider() From c6b6d2510311df7ebcfd9d7009131cbe5176a6ee Mon Sep 17 00:00:00 2001 From: Giuseppe Zileni Date: Wed, 10 Jun 2026 14:33:34 +0200 Subject: [PATCH 2/2] fix(api): drop user-derived prefix from /ingest/json tempfile MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CodeQL flagged the (whitelist-sanitised) safe_stem as "uncontrolled data in path expression". The readable-filename feature was debug ergonomics, not a contract — the URL and source already live inside the file body. Use NamedTemporaryFile's random default name instead. Tests untouched (the contract is the file contents, not the path). --- knowledge-graph-api/api/routes/ingest.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/knowledge-graph-api/api/routes/ingest.py b/knowledge-graph-api/api/routes/ingest.py index e357682..e0b37ad 100644 --- a/knowledge-graph-api/api/routes/ingest.py +++ b/knowledge-graph-api/api/routes/ingest.py @@ -146,10 +146,6 @@ async def ingest_json(body: IngestJsonRequest) -> IngestResult: Returns: An ``IngestResult`` summary. """ - safe_stem = "".join( - c if c.isalnum() or c in "-_ " else "_" for c in body.title - )[:60].strip() or "article" - serialised = ( f"{body.title}\n\n" f"{body.body}\n\n" @@ -159,12 +155,16 @@ async def ingest_json(body: IngestJsonRequest) -> IngestResult: f"Published: {body.published_at.isoformat()}\n" ) + # NamedTemporaryFile defaults: random suffix in the OS temp dir. We + # deliberately do NOT derive the path from request fields — even a + # whitelist-sanitised body field would still trip CodeQL's + # "uncontrolled data in path expression" check, and the readable + # filename is debug-only (the URL/source live in the file body). tmp_path: str | None = None try: with tempfile.NamedTemporaryFile( delete=False, suffix=".txt", - prefix=safe_stem + "_", mode="w", encoding="utf-8", ) as tmp: