refactor: clean code

Author: marcusschiesser

templates/components/services/python/suggestion.py

@@ -24,7 +24,7 @@ def get_configured_prompt(cls) -> Optional[str]:
         return PromptTemplate(prompt)
 
     @classmethod
-    async def suggest_next_questions(
+    async def suggest_next_questions_all_messages(
         cls,
         messages: List[Message],
     ) -> Optional[List[str]]:
@@ -64,3 +64,15 @@ def _extract_questions(cls, text: str) -> List[str]:
         content_match = re.search(r"```(.*?)```", text, re.DOTALL)
         content = content_match.group(1) if content_match else ""
         return content.strip().split("\n")
+
+    @classmethod
+    async def suggest_next_questions(
+        cls,
+        chat_history: List[Message],
+        response: str,
+    ) -> List[str]:
+        """
+        Suggest the next questions that user might ask based on the chat history and the last response
+        """
+        messages = chat_history + [Message(role="assistant", content=response)]
+        return await cls.suggest_next_questions_all_messages(messages)

templates/types/multiagent/fastapi/app/api/routers/vercel_response.py

@@ -1,7 +1,7 @@
 import json
 import logging
 from asyncio import Task
-from typing import AsyncGenerator
+from typing import AsyncGenerator, List
 
 from aiostream import stream
 from app.agents.single import AgentRunEvent, AgentRunResult
@@ -61,38 +61,31 @@ async def _chat_response_generator():
 
             if isinstance(result, AgentRunResult):
                 for token in result.response.message.content:
-                    yield VercelStreamResponse.convert_text(token)
+                    final_response += token
+                    yield cls.convert_text(token)
 
             if isinstance(result, AsyncGenerator):
                 async for token in result:
                     final_response += token.delta
-                    yield VercelStreamResponse.convert_text(token.delta)
+                    yield cls.convert_text(token.delta)
 
             # Generate next questions if next question prompt is configured
-            if NextQuestionSuggestion.get_configured_prompt() is not None:
-                conversation = chat_data.messages + [
-                    Message(role="assistant", content=final_response)
-                ]
-                questions = await NextQuestionSuggestion.suggest_next_questions(
-                    conversation
-                )
-                if questions:
-                    yield VercelStreamResponse.convert_data(
-                        {
-                            "type": "suggested_questions",
-                            "data": questions,
-                        }
-                    )
+            question_data = await cls._generate_next_questions(
+                chat_data.messages, final_response
+            )
+            if question_data:
+                yield cls.convert_data(question_data)
+
             # TODO: stream sources
 
         # Yield the events from the event handler
         async def _event_generator():
             async for event in events():
-                event_response = _event_to_response(event)
+                event_response = cls._event_to_response(event)
                 if verbose:
                     logger.debug(event_response)
                 if event_response is not None:
-                    yield VercelStreamResponse.convert_data(event_response)
+                    yield cls.convert_data(event_response)
 
         combine = stream.merge(_chat_response_generator(), _event_generator())
 
@@ -101,16 +94,28 @@ async def _event_generator():
             if not is_stream_started:
                 is_stream_started = True
                 # Stream a blank message to start the stream
-                yield VercelStreamResponse.convert_text("")
+                yield cls.convert_text("")
 
             async for output in streamer:
                 yield output
                 if await request.is_disconnected():
                     break
 
-
-def _event_to_response(event: AgentRunEvent) -> dict:
-    return {
-        "type": "agent",
-        "data": {"agent": event.name, "text": event.msg},
-    }
+    @staticmethod
+    def _event_to_response(event: AgentRunEvent) -> dict:
+        return {
+            "type": "agent",
+            "data": {"agent": event.name, "text": event.msg},
+        }
+
+    @staticmethod
+    async def _generate_next_questions(chat_history: List[Message], response: str):
+        questions = await NextQuestionSuggestion.suggest_next_questions(
+            chat_history, response
+        )
+        if questions:
+            return {
+                "type": "suggested_questions",
+                "data": questions,
+            }
+        return None

templates/types/multiagent/fastapi/app/api/services/suggestion.py

@@ -26,7 +26,7 @@ class NextQuestions(BaseModel):
 
 class NextQuestionSuggestion:
     @staticmethod
-    async def suggest_next_questions(
+    async def suggest_next_questions_all_messages(
         messages: List[Message],
         number_of_questions: int = N_QUESTION_TO_GENERATE,
     ) -> List[str]:
@@ -58,3 +58,17 @@ async def suggest_next_questions(
         except Exception as e:
             logger.error(f"Error when generating next question: {e}")
             return []
+
+    @staticmethod
+    async def suggest_next_questions(
+        chat_history: List[Message],
+        response: str,
+        number_of_questions: int = N_QUESTION_TO_GENERATE,
+    ) -> List[str]:
+        """
+        Suggest the next questions that user might ask based on the chat history and the last response
+        """
+        messages = chat_history + [Message(role="assistant", content=response)]
+        return await NextQuestionSuggestion.suggest_next_questions_all_messages(
+            messages, number_of_questions
+        )

templates/types/streaming/fastapi/app/api/routers/vercel_response.py

@@ -1,4 +1,5 @@
 import json
+from typing import List
 
 from aiostream import stream
 from fastapi import Request
@@ -54,23 +55,14 @@ async def _chat_response_generator():
             final_response = ""
             async for token in response.async_response_gen():
                 final_response += token
-                yield VercelStreamResponse.convert_text(token)
+                yield cls.convert_text(token)
 
             # Generate next questions if next question prompt is configured
-            if NextQuestionSuggestion.get_configured_prompt() is not None:
-                conversation = chat_data.messages + [
-                    Message(role="assistant", content=final_response)
-                ]
-                questions = await NextQuestionSuggestion.suggest_next_questions(
-                    conversation
-                )
-                if questions:
-                    yield VercelStreamResponse.convert_data(
-                        {
-                            "type": "suggested_questions",
-                            "data": questions,
-                        }
-                    )
+            question_data = await cls._generate_next_questions(
+                chat_data.messages, final_response
+            )
+            if question_data:
+                yield cls.convert_data(question_data)
 
             # the text_generator is the leading stream, once it's finished, also finish the event stream
             event_handler.is_done = True
@@ -93,7 +85,7 @@ async def _event_generator():
             async for event in event_handler.async_event_gen():
                 event_response = event.to_response()
                 if event_response is not None:
-                    yield VercelStreamResponse.convert_data(event_response)
+                    yield cls.convert_data(event_response)
 
         combine = stream.merge(_chat_response_generator(), _event_generator())
         is_stream_started = False
@@ -102,9 +94,21 @@ async def _event_generator():
                 if not is_stream_started:
                     is_stream_started = True
                     # Stream a blank message to start the stream
-                    yield VercelStreamResponse.convert_text("")
+                    yield cls.convert_text("")
 
                 yield output
 
                 if await request.is_disconnected():
                     break
+
+    @staticmethod
+    async def _generate_next_questions(chat_history: List[Message], response: str):
+        questions = await NextQuestionSuggestion.suggest_next_questions(
+            chat_history, response
+        )
+        if questions:
+            return {
+                "type": "suggested_questions",
+                "data": questions,
+            }
+        return None