From a89a1aa80f0af68dbae6bc87281b798eac326fcf Mon Sep 17 00:00:00 2001 From: Niels Provos Date: Mon, 27 Jan 2025 10:50:38 -0800 Subject: [PATCH] fix: update snippet field to be optional in SearchResult and adjust tests accordingly --- src/planai/patterns/search_fetch.py | 6 ++++-- tests/planai/patterns/test_search_fetch.py | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/planai/patterns/search_fetch.py b/src/planai/patterns/search_fetch.py index 939d484..98d42aa 100644 --- a/src/planai/patterns/search_fetch.py +++ b/src/planai/patterns/search_fetch.py @@ -26,7 +26,9 @@ class SearchQuery(Task): class SearchResult(Task): title: str = Field(description="Title of the search result") link: str = Field(description="URL of the search result") - snippet: str = Field(description="Snippet of the search result") + snippet: Optional[str] = Field( + default=None, description="Snippet of the search result" + ) class SearchResults(Task): @@ -54,7 +56,7 @@ class SearchExecutor(CachedTaskWorker): output_types: List[Type[Task]] = [SearchResults] max_results: int = Field(10, description="Maximum number of results per query") - def pre_consume_work(self, task): + def pre_consume_work(self, task: SearchQuery): self.notify_status(task, f"Searching for: {task.query}") def consume_work(self, task: SearchQuery): diff --git a/tests/planai/patterns/test_search_fetch.py b/tests/planai/patterns/test_search_fetch.py index b51d671..36f47bd 100644 --- a/tests/planai/patterns/test_search_fetch.py +++ b/tests/planai/patterns/test_search_fetch.py @@ -40,7 +40,7 @@ def setUp(self): { "title": "Test Result 2", "link": "https://example.com/2", - "snippet": "Test snippet 2", + "snippet": None, }, ]