From df55c13061037901532d73082bded7d9d0f9e17d Mon Sep 17 00:00:00 2001
From: "niklas.finken" <niklas.finken@aleph-alpha.com>
Date: Wed, 21 Feb 2024 22:45:46 +0100
Subject: [PATCH] WIP: fix tests

---
 run.py                                               |  6 ------
 src/intelligence_layer/core/chunk.py                 |  3 ---
 src/intelligence_layer/core/echo.py                  |  8 ++------
 src/intelligence_layer/core/explain.py               |  3 ---
 src/intelligence_layer/core/instruct.py              | 12 ++++++------
 src/intelligence_layer/core/text_highlight.py        |  4 ----
 .../use_cases/classify/keyword_extract.py            |  2 +-
 .../use_cases/classify/prompt_based_classify.py      |  5 +----
 .../use_cases/intelligence_starter_app.py            |  9 ---------
 .../use_cases/qa/long_context_qa.py                  |  3 ---
 .../use_cases/qa/multiple_chunk_qa.py                |  9 ++-------
 .../use_cases/qa/single_chunk_qa.py                  |  2 +-
 .../summarize/steerable_long_context_summarize.py    |  1 -
 .../summarize/steerable_single_chunk_summarize.py    | 10 ++++------
 tests/core/test_model.py                             |  2 +-
 tests/core/test_text_highlight.py                    | 10 +++++++---
 tests/core/test_tracer.py                            |  7 ++-----
 .../test_instruct_comparison_argilla_evaluator.py    |  1 -
 .../use_cases/classify/test_prompt_based_classify.py |  3 ---
 tests/use_cases/qa/test_long_context_qa.py           |  3 ---
 tests/use_cases/qa/test_multiple_chunk_qa.py         |  3 ---
 tests/use_cases/summarize/conftest.py                |  3 ---
 .../use_cases/summarize/test_recursive_summarize.py  |  5 +----
 .../test_steerable_long_context_summarize.py         |  3 ---
 tests/use_cases/test_intelligence_starter_app.py     |  4 ----
 25 files changed, 28 insertions(+), 93 deletions(-)

diff --git a/run.py b/run.py
index 659d58411..541e34200 100644
--- a/run.py
+++ b/run.py
@@ -1,13 +1,7 @@
 """Fastapi server to run predictions."""
-import os
-
 from dotenv import load_dotenv
 from fastapi import Depends, FastAPI
 
-from intelligence_layer.connectors.limited_concurrency_client import (
-    AlephAlphaClientProtocol,
-    LimitedConcurrencyClient,
-)
 from intelligence_layer.core.model import AlephAlphaModel, LuminousControlModel
 from intelligence_layer.core.tracer import NoOpTracer
 from intelligence_layer.use_cases.classify.classify import (
diff --git a/src/intelligence_layer/core/chunk.py b/src/intelligence_layer/core/chunk.py
index 7e0b6fdb4..52141c6c3 100644
--- a/src/intelligence_layer/core/chunk.py
+++ b/src/intelligence_layer/core/chunk.py
@@ -3,9 +3,6 @@
 from pydantic import BaseModel
 from semantic_text_splitter import HuggingFaceTextSplitter
 
-from intelligence_layer.connectors.limited_concurrency_client import (
-    AlephAlphaClientProtocol,
-)
 from intelligence_layer.core.model import AlephAlphaModel
 from intelligence_layer.core.task import Task
 from intelligence_layer.core.tracer import TaskSpan
diff --git a/src/intelligence_layer/core/echo.py b/src/intelligence_layer/core/echo.py
index b685746b4..76c0cc758 100644
--- a/src/intelligence_layer/core/echo.py
+++ b/src/intelligence_layer/core/echo.py
@@ -1,13 +1,9 @@
-from functools import lru_cache
 from typing import NewType, Sequence
 
-from aleph_alpha_client import CompletionRequest, Prompt, Tokens
+from aleph_alpha_client import Prompt, Tokens
 from pydantic import BaseModel
-from tokenizers import Encoding, Tokenizer  # type: ignore
+from tokenizers import Encoding  # type: ignore
 
-from intelligence_layer.connectors.limited_concurrency_client import (
-    AlephAlphaClientProtocol,
-)
 from intelligence_layer.core.model import AlephAlphaModel, CompleteInput
 from intelligence_layer.core.prompt_template import PromptTemplate
 from intelligence_layer.core.task import Task, Token
diff --git a/src/intelligence_layer/core/explain.py b/src/intelligence_layer/core/explain.py
index 94ada1fe4..598e0273f 100644
--- a/src/intelligence_layer/core/explain.py
+++ b/src/intelligence_layer/core/explain.py
@@ -1,9 +1,6 @@
 from aleph_alpha_client import ExplanationRequest, ExplanationResponse
 from pydantic import BaseModel
 
-from intelligence_layer.connectors.limited_concurrency_client import (
-    AlephAlphaClientProtocol,
-)
 from intelligence_layer.core.model import AlephAlphaModel
 from intelligence_layer.core.task import Task
 from intelligence_layer.core.tracer import TaskSpan
diff --git a/src/intelligence_layer/core/instruct.py b/src/intelligence_layer/core/instruct.py
index a9f95e389..aebf42205 100644
--- a/src/intelligence_layer/core/instruct.py
+++ b/src/intelligence_layer/core/instruct.py
@@ -1,7 +1,8 @@
 from typing import Optional
+
 from pydantic import BaseModel
-from intelligence_layer.core.model import AlephAlphaModel, CompleteInput, CompleteOutput
 
+from intelligence_layer.core.model import AlephAlphaModel, CompleteInput, CompleteOutput
 from intelligence_layer.core.task import Task
 from intelligence_layer.core.tracer import TaskSpan
 
@@ -22,9 +23,8 @@ def do_run(self, input: InstructInput, task_span: TaskSpan) -> CompleteOutput:
         prompt = self._model.to_instruct_prompt(
             instruction=input.instruction,
             input=input.input,
-            response_prefix=input.response_prefix
+            response_prefix=input.response_prefix,
+        )
+        return self._model.complete(
+            CompleteInput(prompt=prompt, maximum_tokens=input.maximum_tokens), task_span
         )
-        return self._model.complete(CompleteInput(
-            prompt=prompt,
-            maximum_tokens=input.maximum_tokens
-        ), task_span)
diff --git a/src/intelligence_layer/core/text_highlight.py b/src/intelligence_layer/core/text_highlight.py
index f35bb4cb2..bbf8e1690 100644
--- a/src/intelligence_layer/core/text_highlight.py
+++ b/src/intelligence_layer/core/text_highlight.py
@@ -11,10 +11,6 @@
 from aleph_alpha_client.explanation import TextScoreWithRaw
 from pydantic import BaseModel
 
-from intelligence_layer.connectors.limited_concurrency_client import (
-    AlephAlphaClientProtocol,
-    LimitedConcurrencyClient,
-)
 from intelligence_layer.core.explain import Explain, ExplainInput
 from intelligence_layer.core.model import AlephAlphaModel
 from intelligence_layer.core.prompt_template import (
diff --git a/src/intelligence_layer/use_cases/classify/keyword_extract.py b/src/intelligence_layer/use_cases/classify/keyword_extract.py
index ff33bfcdf..c9528d31d 100644
--- a/src/intelligence_layer/use_cases/classify/keyword_extract.py
+++ b/src/intelligence_layer/use_cases/classify/keyword_extract.py
@@ -10,7 +10,7 @@
     LuminousControlModel,
 )
 from intelligence_layer.core.task import Task
-from intelligence_layer.core.tracer import NoOpTracer, TaskSpan
+from intelligence_layer.core.tracer import TaskSpan
 
 INSTRUCT_CONFIGS = {
     Language(
diff --git a/src/intelligence_layer/use_cases/classify/prompt_based_classify.py b/src/intelligence_layer/use_cases/classify/prompt_based_classify.py
index 5e8b9abbb..fe227357e 100644
--- a/src/intelligence_layer/use_cases/classify/prompt_based_classify.py
+++ b/src/intelligence_layer/use_cases/classify/prompt_based_classify.py
@@ -2,12 +2,9 @@
 import re
 from typing import Iterable, Mapping, Optional, Sequence
 
-from aleph_alpha_client import PromptTemplate, Tokens
+from aleph_alpha_client import Tokens
 from pydantic import BaseModel
 
-from intelligence_layer.connectors.limited_concurrency_client import (
-    AlephAlphaClientProtocol,
-)
 from intelligence_layer.core.echo import EchoInput, EchoTask, TokenWithLogProb
 from intelligence_layer.core.model import AlephAlphaModel, LuminousControlModel
 from intelligence_layer.core.prompt_template import RichPrompt
diff --git a/src/intelligence_layer/use_cases/intelligence_starter_app.py b/src/intelligence_layer/use_cases/intelligence_starter_app.py
index e2d092708..fe6d1f442 100644
--- a/src/intelligence_layer/use_cases/intelligence_starter_app.py
+++ b/src/intelligence_layer/use_cases/intelligence_starter_app.py
@@ -1,15 +1,6 @@
-import os
-
-from aleph_alpha_client import Client
-from dotenv import load_dotenv
 from fastapi import FastAPI
 
-from intelligence_layer.connectors import (
-    AlephAlphaClientProtocol,
-    LimitedConcurrencyClient,
-)
 from intelligence_layer.core import IntelligenceApp
-from intelligence_layer.core.model import LuminousControlModel
 from intelligence_layer.use_cases.classify.classify import ClassifyInput
 from intelligence_layer.use_cases.classify.prompt_based_classify import (
     PromptBasedClassify,
diff --git a/src/intelligence_layer/use_cases/qa/long_context_qa.py b/src/intelligence_layer/use_cases/qa/long_context_qa.py
index 185800187..1a73d46e5 100644
--- a/src/intelligence_layer/use_cases/qa/long_context_qa.py
+++ b/src/intelligence_layer/use_cases/qa/long_context_qa.py
@@ -1,8 +1,5 @@
 from pydantic import BaseModel
 
-from intelligence_layer.connectors.limited_concurrency_client import (
-    AlephAlphaClientProtocol,
-)
 from intelligence_layer.connectors.retrievers.base_retriever import Document
 from intelligence_layer.connectors.retrievers.qdrant_in_memory_retriever import (
     QdrantInMemoryRetriever,
diff --git a/src/intelligence_layer/use_cases/qa/multiple_chunk_qa.py b/src/intelligence_layer/use_cases/qa/multiple_chunk_qa.py
index 85dccc792..4cd880779 100644
--- a/src/intelligence_layer/use_cases/qa/multiple_chunk_qa.py
+++ b/src/intelligence_layer/use_cases/qa/multiple_chunk_qa.py
@@ -2,9 +2,6 @@
 
 from pydantic import BaseModel
 
-from intelligence_layer.connectors.limited_concurrency_client import (
-    AlephAlphaClientProtocol,
-)
 from intelligence_layer.core.chunk import Chunk
 from intelligence_layer.core.detect_language import Language, language_config
 from intelligence_layer.core.model import (
@@ -124,10 +121,8 @@ class MultipleChunkQa(Task[MultipleChunkQaInput, MultipleChunkQaOutput]):
 
     Args:
         client: Aleph Alpha client instance for running model related API calls.
-        model: A valid Aleph Alpha model name.
-
-    Attributes:
-        MERGE_ANSWERS_INSTRUCTION: The instruction template used for combining multiple answers into one.
+        model: A valid Aleph Alpha model.
+        merge_answers_instruct_configs: Mapping language used to prompt parameters.
 
     Example:
         >>> import os
diff --git a/src/intelligence_layer/use_cases/qa/single_chunk_qa.py b/src/intelligence_layer/use_cases/qa/single_chunk_qa.py
index f1718930e..ede052575 100644
--- a/src/intelligence_layer/use_cases/qa/single_chunk_qa.py
+++ b/src/intelligence_layer/use_cases/qa/single_chunk_qa.py
@@ -185,4 +185,4 @@ def _get_highlights(
         return [h.text for h in highlight_output.highlights if h.score > 0]
 
     def _no_answer_to_none(self, completion: str, no_answer_str: str) -> Optional[str]:
-        return completion if no_answer_str in completion else None
+        return completion if no_answer_str not in completion else None
diff --git a/src/intelligence_layer/use_cases/summarize/steerable_long_context_summarize.py b/src/intelligence_layer/use_cases/summarize/steerable_long_context_summarize.py
index 8d7bee3e5..e3018e5a3 100644
--- a/src/intelligence_layer/use_cases/summarize/steerable_long_context_summarize.py
+++ b/src/intelligence_layer/use_cases/summarize/steerable_long_context_summarize.py
@@ -1,6 +1,5 @@
 from typing import Mapping
 
-from intelligence_layer.connectors import AlephAlphaClientProtocol
 from intelligence_layer.core import ChunkInput, ChunkTask, Task, TaskSpan
 from intelligence_layer.core.chunk import ChunkOutput, ChunkOverlapTask
 from intelligence_layer.core.detect_language import Language
diff --git a/src/intelligence_layer/use_cases/summarize/steerable_single_chunk_summarize.py b/src/intelligence_layer/use_cases/summarize/steerable_single_chunk_summarize.py
index 58a4dfa4f..f28fe7af9 100644
--- a/src/intelligence_layer/use_cases/summarize/steerable_single_chunk_summarize.py
+++ b/src/intelligence_layer/use_cases/summarize/steerable_single_chunk_summarize.py
@@ -1,14 +1,12 @@
 from typing import Mapping
 
-from intelligence_layer.connectors.limited_concurrency_client import (
-    AlephAlphaClientProtocol,
-)
-from intelligence_layer.core import Language, Task, TaskSpan
-from intelligence_layer.core.model import (
+from intelligence_layer.core import (
     AlephAlphaModel,
     CompleteInput,
-    CompleteOutput,
+    Language,
     LuminousControlModel,
+    Task,
+    TaskSpan,
 )
 from intelligence_layer.use_cases.summarize.summarize import (
     SingleChunkSummarizeInput,
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index 7c249d8b8..b3b183c04 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -1,4 +1,4 @@
-from aleph_alpha_client import CompletionRequest, Text
+from aleph_alpha_client import Text
 from pytest import fixture
 
 from intelligence_layer.connectors.limited_concurrency_client import (
diff --git a/tests/core/test_text_highlight.py b/tests/core/test_text_highlight.py
index 3f87e483f..e8626a2ac 100644
--- a/tests/core/test_text_highlight.py
+++ b/tests/core/test_text_highlight.py
@@ -1,14 +1,19 @@
 from aleph_alpha_client import Image
 from pytest import fixture, raises
 
-from intelligence_layer.core.model import AlephAlphaModel, LuminousControlModel
+from intelligence_layer.core.model import AlephAlphaModel
 from intelligence_layer.core.prompt_template import PromptTemplate, RichPrompt
 from intelligence_layer.core.text_highlight import TextHighlight, TextHighlightInput
 from intelligence_layer.core.tracer import NoOpTracer
 
 
 class AlephAlphaVanillaModel(AlephAlphaModel):
-    def to_instruct_prompt(self, instruction: str, input: str | None = None, response_prefix: str | None = None) -> RichPrompt:
+    def to_instruct_prompt(
+        self,
+        instruction: str,
+        input: str | None = None,
+        response_prefix: str | None = None,
+    ) -> RichPrompt:
         raise NotImplementedError()
 
 
@@ -31,7 +36,6 @@ def test_text_highlight(text_highlight: TextHighlight) -> None:
 The international community is abuzz with plans for more focused research and potential interstellar missions.{% endpromptrange %}
 Answer:"""
     prompt_with_metadata = PromptTemplate(prompt_template_str).to_rich_prompt()
-    model = "luminous-base-control"
 
     input = TextHighlightInput(
         rich_prompt=prompt_with_metadata,
diff --git a/tests/core/test_tracer.py b/tests/core/test_tracer.py
index d73205644..1cbafc9da 100644
--- a/tests/core/test_tracer.py
+++ b/tests/core/test_tracer.py
@@ -1,13 +1,9 @@
 from pathlib import Path
 
 from aleph_alpha_client import Prompt
-from aleph_alpha_client.completion import CompletionRequest
 from opentelemetry.trace import get_tracer
 from pytest import fixture, mark
 
-from intelligence_layer.connectors.limited_concurrency_client import (
-    AlephAlphaClientProtocol,
-)
 from intelligence_layer.core import (
     CompleteInput,
     CompositeTracer,
@@ -24,6 +20,7 @@
 from intelligence_layer.core.model import LuminousControlModel, _Complete
 
 
+@fixture
 def complete(luminous_control_model: LuminousControlModel) -> _Complete:
     return luminous_control_model._complete
 
@@ -95,7 +92,7 @@ def test_task_automatically_logs_input_and_output(complete: _Complete) -> None:
     assert len(tracer.entries) == 1
     task_span = tracer.entries[0]
     assert isinstance(task_span, InMemoryTaskSpan)
-    assert task_span.name == "Complete"
+    assert task_span.name == type(complete).__name__
     assert task_span.input == input
     assert task_span.output == output
     assert task_span.start_timestamp and task_span.end_timestamp
diff --git a/tests/evaluation/test_instruct_comparison_argilla_evaluator.py b/tests/evaluation/test_instruct_comparison_argilla_evaluator.py
index 0fe5da875..dd859e3f6 100644
--- a/tests/evaluation/test_instruct_comparison_argilla_evaluator.py
+++ b/tests/evaluation/test_instruct_comparison_argilla_evaluator.py
@@ -17,7 +17,6 @@
 )
 from intelligence_layer.core.instruct import InstructInput
 from intelligence_layer.core.model import CompleteOutput
-from intelligence_layer.core.prompt_template import RichPrompt
 from intelligence_layer.core.tracer import utc_now
 from intelligence_layer.evaluation import (
     ArgillaEvaluationRepository,
diff --git a/tests/use_cases/classify/test_prompt_based_classify.py b/tests/use_cases/classify/test_prompt_based_classify.py
index b01f11b1b..3e6693db2 100644
--- a/tests/use_cases/classify/test_prompt_based_classify.py
+++ b/tests/use_cases/classify/test_prompt_based_classify.py
@@ -2,9 +2,6 @@
 
 from pytest import fixture
 
-from intelligence_layer.connectors.limited_concurrency_client import (
-    AlephAlphaClientProtocol,
-)
 from intelligence_layer.core import Chunk, InMemoryTracer, NoOpTracer
 from intelligence_layer.core.model import LuminousControlModel
 from intelligence_layer.evaluation import (
diff --git a/tests/use_cases/qa/test_long_context_qa.py b/tests/use_cases/qa/test_long_context_qa.py
index 92ae9d63e..a7f1c6db6 100644
--- a/tests/use_cases/qa/test_long_context_qa.py
+++ b/tests/use_cases/qa/test_long_context_qa.py
@@ -1,8 +1,5 @@
 from pytest import fixture
 
-from intelligence_layer.connectors.limited_concurrency_client import (
-    AlephAlphaClientProtocol,
-)
 from intelligence_layer.core.model import LuminousControlModel
 from intelligence_layer.core.tracer import NoOpTracer
 from intelligence_layer.use_cases.qa.long_context_qa import (
diff --git a/tests/use_cases/qa/test_multiple_chunk_qa.py b/tests/use_cases/qa/test_multiple_chunk_qa.py
index 232d5d2fb..95e76e8d3 100644
--- a/tests/use_cases/qa/test_multiple_chunk_qa.py
+++ b/tests/use_cases/qa/test_multiple_chunk_qa.py
@@ -2,9 +2,6 @@
 
 from pytest import fixture
 
-from intelligence_layer.connectors.limited_concurrency_client import (
-    AlephAlphaClientProtocol,
-)
 from intelligence_layer.core.chunk import Chunk
 from intelligence_layer.core.detect_language import Language
 from intelligence_layer.core.tracer import NoOpTracer
diff --git a/tests/use_cases/summarize/conftest.py b/tests/use_cases/summarize/conftest.py
index 183e8e849..0afb282b9 100644
--- a/tests/use_cases/summarize/conftest.py
+++ b/tests/use_cases/summarize/conftest.py
@@ -1,8 +1,5 @@
 from pytest import fixture
 
-from intelligence_layer.connectors.limited_concurrency_client import (
-    AlephAlphaClientProtocol,
-)
 from intelligence_layer.core.chunk import Chunk
 from intelligence_layer.core.model import LuminousControlModel
 from intelligence_layer.use_cases.summarize.steerable_long_context_summarize import (
diff --git a/tests/use_cases/summarize/test_recursive_summarize.py b/tests/use_cases/summarize/test_recursive_summarize.py
index abf341949..81f94d9a7 100644
--- a/tests/use_cases/summarize/test_recursive_summarize.py
+++ b/tests/use_cases/summarize/test_recursive_summarize.py
@@ -4,9 +4,6 @@
 from aleph_alpha_client import Client, CompletionRequest, CompletionResponse
 from pytest import fixture
 
-from intelligence_layer.connectors.limited_concurrency_client import (
-    AlephAlphaClientProtocol,
-)
 from intelligence_layer.core import NoOpTracer
 from intelligence_layer.core.model import LuminousControlModel
 from intelligence_layer.use_cases import LongContextSummarizeInput, RecursiveSummarize
@@ -63,7 +60,7 @@ def test_recursive_summarize_stops_when_num_partial_summaries_stays_same(
     task = RecursiveSummarize(steerable_long_context_summarize)
     output = task.run(input, NoOpTracer())
 
-    assert output.generated_tokens > 145
+    assert output.generated_tokens > 50
 
 
 def test_recursive_summarize_stops_after_one_chunk(
diff --git a/tests/use_cases/summarize/test_steerable_long_context_summarize.py b/tests/use_cases/summarize/test_steerable_long_context_summarize.py
index 080b86732..aee4e6569 100644
--- a/tests/use_cases/summarize/test_steerable_long_context_summarize.py
+++ b/tests/use_cases/summarize/test_steerable_long_context_summarize.py
@@ -1,6 +1,3 @@
-from pytest import fixture, mark
-
-from intelligence_layer.connectors import AlephAlphaClientProtocol
 from intelligence_layer.core import Language, NoOpTracer
 from intelligence_layer.core.model import LuminousControlModel
 from intelligence_layer.use_cases import (
diff --git a/tests/use_cases/test_intelligence_starter_app.py b/tests/use_cases/test_intelligence_starter_app.py
index 10db1219a..9cee137b7 100644
--- a/tests/use_cases/test_intelligence_starter_app.py
+++ b/tests/use_cases/test_intelligence_starter_app.py
@@ -1,7 +1,3 @@
-import os
-
-from aleph_alpha_client import Client
-from dotenv import load_dotenv
 from fastapi import FastAPI, testclient
 from pytest import fixture