fix: update snippet field to be optional in SearchResult and adjust t…

…ests accordingly
provos · Jan 27, 2025 · a89a1aa · a89a1aa
1 parent 1957f48
commit a89a1aa
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git 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
@@ -40,7 +40,7 @@ def setUp(self):
             {
                 "title": "Test Result 2",
                 "link": "https://example.com/2",
-                "snippet": "Test snippet 2",
+                "snippet": None,
             },
         ]