Merge pull request #226 from restackio/bumpComunnityDefense

bump and add pydantic

Author: aboutphilippe

community/defense_quickstart_news_scraper_summarizer/pyproject.toml

@@ -7,7 +7,7 @@ requires-python = ">=3.10,<3.14"
 readme = "README.md"
 dependencies = [
     "pydantic>=2.10.6",
-    "restack-ai==0.0.81",
+    "restack-ai==0.0.87",
     "watchfiles>=1.0.4",
     "python-dotenv==1.0.1",
     "openai>=1.61.0",

community/defense_quickstart_news_scraper_summarizer/src/functions/crawl/website.py

@@ -1,12 +1,16 @@
 from restack_ai.function import function, log
 import requests
 from bs4 import BeautifulSoup
+from pydantic import BaseModel
+
+class CrawlInput(BaseModel):
+    url: str
 
 @function.defn()
-async def crawl_website(url):
+async def crawl_website(input: CrawlInput):
     try:
         # Send a GET request to the URL
-        response = requests.get(url)
+        response = requests.get(input.url)
         response.raise_for_status()  # Raise an error for bad responses
 
         # Parse the content with BeautifulSoup

community/defense_quickstart_news_scraper_summarizer/src/functions/helper/split_text.py

@@ -1,16 +1,23 @@
 from restack_ai.function import function
 
+from pydantic import BaseModel
+
+class SplitTextInput(BaseModel):
+    text: str
+    average_token_per_character: int = 3
+    max_tokens: int = 4096
+
 @function.defn()
-async def split_text(text: str, average_token_per_character: int = 3, max_tokens: int = 4096) -> list:
+async def split_text(input: SplitTextInput) -> list:
     chunks = []
     current_chunk = []
     current_length = 0
 
-    for char in text:
+    for char in input.text:
         current_chunk.append(char)
-        current_length += average_token_per_character
+        current_length += input.average_token_per_character
 
-        if current_length >= max_tokens:
+        if current_length >= input.max_tokens:
             chunks.append(''.join(current_chunk))
             current_chunk = []
             current_length = 0

community/defense_quickstart_news_scraper_summarizer/src/workflows/workflow.py

@@ -3,19 +3,28 @@
 
 with import_functions():
     from src.functions.rss.pull import rss_pull
-    from src.functions.crawl.website import crawl_website
-    from src.functions.helper.split_text import split_text
+    from src.functions.crawl.website import crawl_website, CrawlInput
+    from src.functions.helper.split_text import split_text, SplitTextInput
     from src.functions.llm.chat import llm_chat, FunctionInputParams
     from src.functions.rss.schema import RssInput
 
+from pydantic import BaseModel
+
+class RssWorkflowInput(BaseModel):
+    url: str
+    count: int
+
 @workflow.defn()
 class RssWorkflow:
     @workflow.run
-    async def run(self, input: dict):
+    async def run(self, input: RssWorkflowInput):
 
-        url = input["url"]
-        count = input["count"]
-        rss_results = await workflow.step(rss_pull, RssInput(url=url, count=count), start_to_close_timeout=timedelta(seconds=10))
+        url = input.url
+        count = input.count
+        rss_results = await workflow.step(
+            function=rss_pull,
+            function_input=RssInput(url=url, count=count),
+            start_to_close_timeout=timedelta(seconds=10))
         urls = [item['link'] for item in rss_results if 'link' in item]
         titles = [item['title'] for item in rss_results if 'title' in item]
 
@@ -25,8 +34,14 @@ async def run(self, input: dict):
             log.info("rss_result", extra={"url": url})
             if url:
                 try:
-                    content = await workflow.step(crawl_website, url, start_to_close_timeout=timedelta(seconds=30))
-                    split_content = await workflow.step(split_text, f"{titles[urls.index(url)]}\n\n{content}", start_to_close_timeout=timedelta(seconds=30))
+                    content = await workflow.step(
+                        function=crawl_website,
+                        function_input=CrawlInput(url),
+                        start_to_close_timeout=timedelta(seconds=30))
+                    split_content = await workflow.step(
+                        function=split_text,
+                        function_input=SplitTextInput(text=f"{titles[urls.index(url)]}\n\n{content}"),
+                        start_to_close_timeout=timedelta(seconds=30))
                     crawled_contents.append(split_content)
                 except Exception as e:
                     log.error(f"Failed to crawl {url}: {str(e)}")
@@ -35,12 +50,24 @@ async def run(self, input: dict):
         for split_content in crawled_contents:
             for content in split_content:
                 user_prompt = f"Provide a translation of the news article. Translate the following content to English: {content}"
-                translation = await workflow.step(llm_chat, FunctionInputParams(user_prompt=user_prompt), task_queue="llm_chat",start_to_close_timeout=timedelta(seconds=120))
+                translation = await workflow.step(
+                    function=llm_chat,
+                    function_input=FunctionInputParams(user_prompt=user_prompt),
+                    task_queue="llm_chat",
+                    start_to_close_timeout=timedelta(seconds=120))
 
                 user_prompt = f"Provide a summary of the news found on rss feed. Summarize the following content: {translation} in maxium 1 sentence with no more than 20 words"
-                summary = await workflow.step(llm_chat, FunctionInputParams(user_prompt=user_prompt), task_queue="llm_chat",start_to_close_timeout=timedelta(seconds=120))
+                summary = await workflow.step(
+                    function=llm_chat,
+                    function_input=FunctionInputParams(user_prompt=user_prompt),
+                    task_queue="llm_chat",start_to_close_timeout=timedelta(seconds=120))
                 summaries.append(summary)
 
         user_prompt = f"Make a daily digest of all the news and tell me what is the most important news. Here are the summaries of the articles: {summaries}."
 
-        return await workflow.step(llm_chat, FunctionInputParams(user_prompt=user_prompt), task_queue="llm_chat", start_to_close_timeout=timedelta(seconds=120))
+        return await workflow.step(
+            function=llm_chat, 
+            funcion_input=FunctionInputParams(user_prompt=user_prompt),
+            task_queue="llm_chat",
+            start_to_close_timeout=timedelta(seconds=120)
+        )