mypy + tests

Signed-off-by: Henry Lindeman <[email protected]>
aryn-ai · Feb 27, 2025 · d156b03 · d156b03
1 parent 5b267a7
commit d156b03
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 8 deletions.
diff --git a/lib/sycamore/sycamore/query/execution/operations.py b/lib/sycamore/sycamore/query/execution/operations.py
@@ -1,5 +1,5 @@
 import math
-from typing import Any, List, Union, Optional
+from typing import Any, List, Union, Optional, Type
 
 import structlog
 
@@ -21,14 +21,15 @@
 
 log = structlog.get_logger(__name__)
 # multistep
-DEFAULT_DOCSET_SUMMARIZER_CLS = MultiStepDocumentSummarizer
+DEFAULT_DOCSET_SUMMARIZER_CLS = MultiStepDocumentSummarizer  # type: ignore
+
 DEFAULT_SUMMARIZER_KWARGS: dict[str, Any] = {
     "fields": "*",
     "tokenizer": OpenAITokenizer("gpt-4o"),
     "max_tokens": 80_000,
 }
 # onestep
-DEFAULT_DOCSET_SUMMARIZER_CLS = OneStepDocumentSummarizer
+DEFAULT_DOCSET_SUMMARIZER_CLS = OneStepDocumentSummarizer  # type: ignore
 DEFAULT_SUMMARIZER_KWARGS = {"fields": [EtCetera], "tokenizer": OpenAITokenizer("gpt-4o"), "token_limit": 80_000}
 
 
@@ -92,7 +93,7 @@ def summarize_data(
         Conversational response to question.
     """
     if docset_summarizer is None:
-        docset_summarizer = DEFAULT_DOCSET_SUMMARIZER_CLS(llm=llm, question=question, **DEFAULT_SUMMARIZER_KWARGS)
+        docset_summarizer = DEFAULT_DOCSET_SUMMARIZER_CLS(llm=llm, question=question, **DEFAULT_SUMMARIZER_KWARGS)  # type: ignore
 
     if all(isinstance(d, DocSet) for d in result_data):
         return summarize_data_docsets(

diff --git a/lib/sycamore/sycamore/tests/unit/query/test_operations.py b/lib/sycamore/sycamore/tests/unit/query/test_operations.py
@@ -16,7 +16,7 @@
     summarize_data,
     math_operation,
 )
-from sycamore.transforms.summarize import NUM_DOCS_GENERATE
+from sycamore.transforms.summarize import NUM_DOCS_GENERATE, MultiStepDocumentSummarizer
 
 
 class MockLLM(LLM):
@@ -143,7 +143,13 @@ def test_summarize_data(self, words_and_ids_docset):
     def test_get_text_for_summarize_data_docset(self, words_and_ids_docset):
         llm = MockLLM()
         summarize_data(
-            llm=llm, question=None, result_description="List of unique cities", result_data=[words_and_ids_docset]
+            llm=llm,
+            question=None,
+            result_description="List of unique cities",
+            result_data=[words_and_ids_docset],
+            docset_summarizer=MultiStepDocumentSummarizer(
+                llm=llm, question=None, data_description="List of unique cities"
+            ),
         )
         captured = llm.capture[-1]
         mcontent = captured.messages[-1].content
@@ -160,9 +166,12 @@ def test_get_text_for_summarize_data_docset_with_elements(self, big_words_and_id
             result_description="List of unique cities",
             result_data=[big_words_and_ids_docset],
             summaries_as_text=True,
+            docset_summarizer=MultiStepDocumentSummarizer(
+                llm=llm, question=None, data_description="List of unique cities", max_tokens=1000
+            ),
         )
         captured = llm.capture
-        assert len(captured) == 45  # 45 llm calls
+        assert len(captured) == 44  # 44 llm calls
         assert response == "merged summary"
 
     def test_get_text_for_summarize_data_non_docset(self, words_and_ids_docset):

diff --git a/lib/sycamore/sycamore/transforms/summarize.py b/lib/sycamore/sycamore/transforms/summarize.py
@@ -184,6 +184,7 @@ def __init__(
         self,
         llm: LLM,
         question: Optional[str] = None,
+        data_description: Optional[str] = None,
         prompt: SycamorePrompt = MaxTokensHeirarchyPrompt,
         fields: Union[None, Literal["*"], list[str]] = None,
         max_tokens: int = 10 * 1000,
@@ -194,6 +195,7 @@ def __init__(
         self.prompt = prompt.set(**self.get_const_vars())
         self.fields = fields
         self.question = question
+        self.data_description = data_description
         self.max_tokens = max_tokens
         self.tokenizer = tokenizer
         self.rounds = 4
@@ -246,6 +248,8 @@ def as_llm_map(self, child: Optional[Node], **kwargs) -> Node:
             self.prompt = self.prompt.set(fields=self.fields)
         if self.question is not None:
             self.prompt = self.prompt.set(question=self.question)
+        if self.data_description is not None:
+            self.prompt = self.prompt.set(data_description=self.data_description)
         nodes = []
         last = child
         for round in range(self.rounds):
@@ -361,7 +365,6 @@ def preprocess(self, doc: Document) -> Document:
         this = self.prompt.render_document(doc)
         while last != this:
             ntk = this.token_count(self.tokenizer)
-            print(ntk)
             if ntk > self.token_limit:
                 doc.properties[vars["numel_key"]] -= 1
                 return doc