diff --git a/example_segmentation.py b/example_segmentation.py index 96ef1f9..f4b141c 100644 --- a/example_segmentation.py +++ b/example_segmentation.py @@ -6,16 +6,42 @@ import openai class Search(OpenAISchema): - "Search query for a request," + """ + Search query for a single request + + Tips: + - Be specific with your query, use key words and multiple representations of the same thing, e.g. "video" and "video clip" or "SSO" and "single sign on" + - Use the title to describe the request, e.g. "Video from last week about the investment case study" + """ title: str = Field(..., description="Title of the request") query: str = Field(..., description="Query to search for relevant content") + async def execute(self): + import asyncio + + await asyncio.sleep(1) + print(f"Searching for `{self.title}` with query `{self.query}`") + class MultiSearch(OpenAISchema): - "Search query for multiple requests" + """ + Segment a request into multiple search queries + + Tips: + - Do not overlap queries, e.g. "video" and "video clip" are too similar + """ + searches: List[Search] = Field(..., description="List of searches") + def execute(self): + import asyncio + + loop = asyncio.get_event_loop() + + tasks = asyncio.gather(*[search.execute() for search in self.searches]) + return loop.run_until_complete(tasks) + @retry(stop=stop_after_attempt(3)) def segment(data: str) -> MultiSearch: @@ -42,5 +68,5 @@ if __name__ == "__main__": queries = segment( "Please send me the video from last week about the investment case study and also documents about your GPDR policy?" ) - for query in queries.searches: - print(query) + + results = queries.execute()