diff --git a/tests/test_api.py b/tests/test_api.py index 4c57f38..f726866 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -17,7 +17,11 @@ def handler(request: httpx2.Request) -> httpx2.Response: ) mock_http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler)) - generator = ClaudeAnswerGenerator(api_key="test-key", model="test-model", client=mock_http_client) + generator = ClaudeAnswerGenerator( + api_key="test-key", + model="test-model", + client=mock_http_client, + ) pipeline = RagPipeline(generator=generator) app = create_app(pipeline=pipeline) return TestClient(app) @@ -40,7 +44,10 @@ def test_querying_before_any_document_is_ingested_returns_no_sources(client): def test_querying_after_ingesting_returns_a_cited_answer(client): - client.post("/documents", json={"document_id": "animals", "text": "Cats are small domesticated mammals."}) + client.post( + "/documents", + json={"document_id": "animals", "text": "Cats are small domesticated mammals."}, + ) response = client.post("/query", json={"question": "Are cats mammals?"}) diff --git a/tests/test_generation.py b/tests/test_generation.py index 889fdcf..fd9ff30 100644 --- a/tests/test_generation.py +++ b/tests/test_generation.py @@ -11,8 +11,14 @@ def test_grounded_prompt_includes_every_source_and_the_question(): results = [ - SearchResult(chunk=Chunk(document_id="doc-a", index=0, text="Cats are mammals."), score=0.9), - SearchResult(chunk=Chunk(document_id="doc-b", index=1, text="Dogs are mammals too."), score=0.5), + SearchResult( + chunk=Chunk(document_id="doc-a", index=0, text="Cats are mammals."), + score=0.9, + ), + SearchResult( + chunk=Chunk(document_id="doc-b", index=1, text="Dogs are mammals too."), + score=0.5, + ), ] prompt = build_grounded_prompt("Are cats mammals?", results) @@ -26,12 +32,20 @@ def test_grounded_prompt_includes_every_source_and_the_question(): async def test_generator_sends_the_grounded_prompt_and_returns_the_answer(build_mock_client): captured = {} - results = [SearchResult(chunk=Chunk(document_id="doc-a", index=0, text="Cats are mammals."), score=0.9)] + results = [ + SearchResult( + chunk=Chunk(document_id="doc-a", index=0, text="Cats are mammals."), + score=0.9, + ), + ] def handler(request: httpx2.Request) -> httpx2.Response: captured["headers"] = request.headers captured["body"] = json.loads(request.content) - return httpx2.Response(200, json={"content": [{"type": "text", "text": "Yes, cats are mammals [source 1]."}]}) + return httpx2.Response( + 200, + json={"content": [{"type": "text", "text": "Yes, cats are mammals [source 1]."}]}, + ) generator = ClaudeAnswerGenerator( api_key="test-key", model="test-model", client=build_mock_client(handler) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index ea62e2f..b27bc92 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -31,7 +31,10 @@ async def test_ingest_reports_how_many_chunks_it_created(): async def test_ask_returns_the_generated_answer_with_matching_sources(): def handler(request: httpx2.Request) -> httpx2.Response: - return httpx2.Response(200, json={"content": [{"type": "text", "text": "Cats are mammals."}]}) + return httpx2.Response( + 200, + json={"content": [{"type": "text", "text": "Cats are mammals."}]}, + ) pipeline = _pipeline_with_generator(handler) pipeline.ingest("animals", "Cats are small domesticated feline mammals.")