feat: add Premium feature

src/agents_framework/integrations.py

@@ -17,8 +17,9 @@ from .tools import (
     generate_response_base_context,
     openai_client_answer
 )
+from api.schemas import OutputLocation
 from .utils import get_data
-from typing import List
+from typing import List, Dict, Tuple, Union
 
 # log = logger.get_logger(__name__)
 
@@ -55,27 +56,26 @@ class AgentIntegrationService:
             return_source_documents=True,
         )
 
-    async def answer_to_messages(self, query: List[list], chat_history: List[dict] | None) -> str:
-        question_type = classifer_question_execution.invoke(query).content.replace('`', "").replace("Tags", "").strip()
-        # log.info(f"question_type: {question_type}")
-        print(question_type)
+    async def answer_to_messages(self, query: List[list], chat_history: List[dict] | None, question_type: str) -> Union[str, OutputLocation | None]:
         agent_using = self._route(question_type)
         if agent_using == self.answer_based_search:
-            query_reformat = openai_client_answer(query, chat_history, model="gpt-3.5-turbo", template=data["rephrase_query"])
-            if query == 'not_need':
-                query_reformat = query
-            print(query_reformat)
+            query_reformat = openai_client_answer(
+                query, 
+                chat_history, 
+                model="gpt-3.5-turbo", 
+                template=data["rephrase_query"]
+            )
+            if query == 'not_need': query_reformat = query
             answer = agent_using(query_reformat)
         else:
             answer = agent_using(query, chat_history)
 
-        if question_type in ["ask_personal", "search_web"]:
-            return {"text": answer.content, "ids_location": None}
+        text_message, location_message = "", None
+        if question_type in ["ask_personal", "search_web"]: text_message = answer
         else:
-            return {
-                "text": f"Đây là lộ trình tôi nghĩ ra bao gồm {len(set(answer))} điểm du lịch", 
-                "ids_location": list(set(answer))
-            }
+            text_message = f"Chào bạn, đây là một lịch trình tham quan thành phố Hồ Chí Minh trong một ngày mà bạn có thể tham khảo:"
+            location_message = OutputLocation(locations=answer, message=text_message)
+        return text_message, location_message
 
     
     def answer_based_retrieval(self, query, chat_history):
@@ -93,16 +93,16 @@ class AgentIntegrationService:
             else:
                 continue
         num_questions, questions = batch["number_of_locations"], batch["questions"]
-        locations_index = []
+        locations_index = set()
         print(questions)
         for question in questions:
-            locations_index.append(self._retrieval_places(question))
-        return locations_index
+            locations_index.add(self._retrieval_places(question))
+        return list(locations_index)
+
 
     def _retrieval_places(self, query):
         response = self.qa_baseline.invoke({"query": query})
         location_index = response["source_documents"][0].metadata["seq_num"] - 1
-        print(location_index)
         return location_index
     
     
@@ -114,8 +114,8 @@ class AgentIntegrationService:
         date = datetime.datetime.now().isoformat()
         print(context)
         answer = generate_response_base_context(context.strip(), date=date, template=data['interpreter_web'])
-        print(answer)
-        return answer
+        # print(answer.content)
+        return answer.content
 
     
     def answer_normal(self, user_input: str, chat_history, model_name="gpt-3.5-turbo") -> str:

src/api/routes.py

@@ -3,7 +3,8 @@ from fastapi import APIRouter, HTTPException, Depends
 from fastapi.security import APIKeyHeader
 
 from agents.models import service
-from .schemas import MessageInput ,ChatAgentResponse
+from agents_framework.tools import classifer_question_execution
+from .schemas import MessageInput ,ChatAgentResponse, OutputLocation
 
 
 header_scheme = APIKeyHeader(name="X-API-Key", auto_error=False)
@@ -25,6 +26,7 @@ async def agents_root():
 @router.post("/chat-agent")
 async def chat_completion(
     message: MessageInput,
+    user_status,
     api_key: APIKeyHeader = Depends(header_scheme)
                     
 ) -> ChatAgentResponse:
@@ -53,23 +55,26 @@ async def chat_completion(
     user_question = message["question"]
     hist_messages = message["history"]
 
-    # log.info(f"User question: {user_question}")
-    # log.info(f"User message: {hist_messages}")
-
     chat_history = []
     if hist_messages and len(hist_messages[0]) > 0:
         chat_history = hist_messages
-        # for mes in hist_messages:
-            # log.info(f"Conversation history message: {mes['question']} | {mes['answer']}")
-            # chat_history.append(HumanMessage(content=mes['question']))
-            # chat_history.append(AIMessage(content=mes['answer']))
     
-    # log.info(f"User Question: {user_question}")
+    question_type = classify_question_type(user_question)
+    if question_type == "search_web" and user_status != "Premium":
+        return ChatAgentResponse(texts_message="Rất tiếc, hiện tại bạn chưa thể sử dụng tính năng này. Hãy nâng cấp lên gói Premium ngay hôm nay để có thể trải nghiệm không giới hạn và mở khóa những tính năng tuyệt vời khác cho chuyến đ cho mình!")
 
-    response = await service.answer_to_messages(query=user_question, chat_history=chat_history)
+    text_message, location_message = await service.answer_to_messages(query=user_question, chat_history=chat_history, question_type=question_type)
     # log.info(f"Agent response: {response['text']}, Sources: {response['ids_location']}")
 
-    return response
+    return ChatAgentResponse(texts_message=text_message, locations_message=location_message) 
+    # return {"text": text, "ids_location": ids_location}
+
+def classify_question_type(query: str): 
+
+    question_type = classifer_question_execution.invoke(query).content.replace('`', "").replace("Tags", "").strip()
+    # log.info(f"question_type: {question_type}")
+    print(question_type)
+    return question_type
 
 
 

src/api/schemas.py

@@ -2,7 +2,7 @@ from typing import List
 from pydantic import BaseModel
 
 from langchain.docstore.document import Document
-
+from typing import Optional
 
 # # #########################################
 # # Internal schemas
@@ -17,6 +17,11 @@ class MessageInput(BaseModel):
 # # API schemas
 # # #########################################
 
+class OutputLocation(BaseModel):
+    locations: List[int]
+    message: str
+
 class ChatAgentResponse(BaseModel):
-    text: str | None
-    ids_location: List[int] | None
+    texts_message: str | None
+    locations_message: Optional[None | OutputLocation] = None
+