Enable GPT-4o image feature (for API only)

Author: chatrtham

src/fastapi_app/api_models.py

@@ -1,15 +1,25 @@
-from typing import Any
+from typing import Any, List, Union, Optional
 
 from pydantic import BaseModel
 
+class TextContent(BaseModel):
+    type: str
+    text: str
+
+class ImageUrl(BaseModel):
+    url: str
+    detail: str = "auto"
+
+class ImageContent(BaseModel):
+    type: str
+    image_url: ImageUrl
 
 class Message(BaseModel):
-    content: str
     role: str = "user"
-
+    content: Union[str, List[Union[TextContent, ImageContent]]]
 
 class ChatRequest(BaseModel):
-    messages: list[Message]
+    messages: List[Message]
     context: dict = {}
 
 

src/fastapi_app/postgres_models.py

@@ -14,7 +14,7 @@ class Base(DeclarativeBase, MappedAsDataclass):
 
 class Item(Base):
     __tablename__ = "packages"
-    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
+    id: Mapped[int] = mapped_column(primary_key=True)
     package_name: Mapped[str] = mapped_column()
     package_picture: Mapped[str] = mapped_column()
     url: Mapped[str] = mapped_column()
@@ -59,7 +59,7 @@ class Item(Base):
     embedding_url: Mapped[Vector] = mapped_column(Vector(1536))
     embedding_installment_month: Mapped[Vector] = mapped_column(Vector(1536))
     embedding_installment_limit: Mapped[Vector] = mapped_column(Vector(1536))
-    embedding_price_to_reserve_for_this_package: Mapped[Vector] = mapped_column(Vector(1536))
+    embedding_price_to_reserve_for_this_package: Mapped[Vector] = (mapped_column(Vector(1536)))
     embedding_shop_name: Mapped[Vector] = mapped_column(Vector(1536))
     embedding_category: Mapped[Vector] = mapped_column(Vector(1536))
     embedding_category_tags: Mapped[Vector] = mapped_column(Vector(1536))
@@ -178,7 +178,16 @@ def to_dict(self, include_embedding: bool = False):
                 del model_dict[col]
         return model_dict
 
-    def to_str_for_rag(self):
+    def to_str_for_broad_rag(self):
+        return f"""
+    package_name: {self.package_name}
+    url: {self.url}
+    locations: {self.locations}
+    price: {self.price}
+    brand: {self.brand}
+    """
+
+    def to_str_for_narrow_rag(self):
         return f"""
     package_name: {self.package_name}
     package_picture: {self.package_picture}

src/fastapi_app/prompts/answer.txt

@@ -2,7 +2,8 @@ You are a woman name Jib (จิ๊บ), a sales consultant from HDmall which yo
 Represent youself as จิ๊บ and the customer as คุณลูกค้า when the conversation is in Thai language.
 Represent youself as Jib and the customer as you when the conversation is in English language.
 When the customer asks about any packages, please make sure to provide brand, price , URL and location every time.
-For EVERY URL that you provided, create a new line then wrap it with <a href=""> so the user can open it via UI
 If customer wants to talk with admin, please provide the URL link in the datasource.
+Note that some packages may have additional cost.
 If the user is asking a question regarding location, proximity, or area, query relevant document from the source and ask where the user is at. please try to suggest services or answers closest to the location the user is asking as much as possible.
 Answer ONLY with the facts listed in the list of sources below. If there isn't enough information below, say sorry you don't know. Do not generate answers that don't use the sources below. If asking a clarifying question to the user would help, ask the question.
+Answer in well-structured plain text, not a markdown.

src/fastapi_app/query_rewriter.py

@@ -18,38 +18,41 @@ def build_search_function() -> list[ChatCompletionToolParam]:
                     "properties": {
                         "search_query": {
                             "type": "string",
-                            "description": "Query string to use for full text search, e.g. 'red shoes'",
+                            "description": "Query string to use for full text search, e.g. 'ตรวจสุขภาพ'",
                         },
                         "price_filter": {
                             "type": "object",
-                            "description": "Filter search results based on price of the product",
+                            "description": "Filter search results based on price in Thai Baht of the package",
                             "properties": {
                                 "comparison_operator": {
                                     "type": "string",
-                                    "description": "Operator to compare the column value, either '>', '<', '>=', '<=', '=='",  # noqa
+                                    "description": "Operator to compare the column value, either '>', '<', '>=', '<=', '='",  # noqa
                                 },
                                 "value": {
                                     "type": "number",
                                     "description": "Value to compare against, e.g. 30",
                                 },
                             },
                         },
-                        # "brand_filter": {
-                        #     "type": "object",
-                        #     "description": "Filter search results based on brand of the product",
-                        #     "properties": {
-                        #         "comparison_operator": {
-                        #             "type": "string",
-                        #             "description": "Operator to compare the column value, either '==' or '!='",
-                        #         },
-                        #         "value": {
-                        #             "type": "string",
-                        #             "description": "Value to compare against, e.g. AirStrider",
-                        #         },
-                        #     },
-                        # },
+                        "url_filter": {
+                            "type": "object",
+                            "description": "Filter search results based on url of the package. The url is package specific.",
+                            "properties": {
+                                "comparison_operator": {
+                                    "type": "string",
+                                    "description": "Operator to compare the column value, either '=' or '!='",
+                                },
+                                "value": {
+                                    "type": "string",
+                                    "description": """
+                                    The package URL to compare against.
+                                    Don't pass anything if you can't specify the exact URL from user query.
+                                    """,
+                                },
+                            },
+                        },
                     },
-                    "required": ["search_query"],
+                    "required": ["search_query", "url_filter"],
                 },
             },
         }
@@ -67,6 +70,7 @@ def extract_search_arguments(chat_completion: ChatCompletion):
             function = tool.function
             if function.name == "search_database":
                 arg = json.loads(function.arguments)
+                print(arg)
                 search_query = arg.get("search_query")
                 if "price_filter" in arg and arg["price_filter"]:
                     price_filter = arg["price_filter"]
@@ -77,15 +81,16 @@ def extract_search_arguments(chat_completion: ChatCompletion):
                             "value": price_filter["value"],
                         }
                     )
-                # if "brand_filter" in arg and arg["brand_filter"]:
-                #     brand_filter = arg["brand_filter"]
-                #     filters.append(
-                #         {
-                #             "column": "brand",
-                #             "comparison_operator": brand_filter["comparison_operator"],
-                #             "value": brand_filter["value"],
-                #         }
-                #     )
+                if "url_filter" in arg and arg["url_filter"]:
+                    url_filter = arg["url_filter"]
+                    if url_filter["value"] != "https://hdmall.co.th":
+                        filters.append(
+                            {
+                                "column": "url",
+                                "comparison_operator": url_filter["comparison_operator"],
+                                "value": url_filter["value"],
+                            }
+                        )
     elif query_text := response_message.content:
         search_query = query_text.strip()
     return search_query, filters

src/fastapi_app/rag_advanced.py

@@ -1,3 +1,4 @@
+import copy
 import pathlib
 from collections.abc import AsyncGenerator
 from typing import (
@@ -8,7 +9,7 @@
 from openai.types.chat import (
     ChatCompletion,
 )
-from openai_messages_token_helper import build_messages, get_token_limit
+from openai_messages_token_helper import get_token_limit
 
 from .api_models import ThoughtStep
 from .postgres_searcher import PostgresSearcher
@@ -36,23 +37,21 @@ def __init__(
     async def run(
         self, messages: list[dict], overrides: dict[str, Any] = {}
     ) -> dict[str, Any] | AsyncGenerator[dict[str, Any], None]:
+        # Normalize the message format
+        for message in messages:
+            if isinstance(message['content'], str):
+                message['content'] = [{'type': 'text', 'text': message['content']}]
+
+        # Determine the search mode and the number of results to return
         text_search = overrides.get("retrieval_mode") in ["text", "hybrid", None]
         vector_search = overrides.get("retrieval_mode") in ["vectors", "hybrid", None]
         top = overrides.get("top", 3)
 
-        original_user_query = messages[-1]["content"]
-        past_messages = messages[:-1]
 
         # Generate an optimized keyword search query based on the chat history and the last question
+        query_messages = copy.deepcopy(messages)
+        query_messages.insert(0, {"role": "system", "content": self.query_prompt_template})
         query_response_token_limit = 500
-        query_messages = build_messages(
-            model=self.chat_model,
-            system_prompt=self.query_prompt_template,
-            new_user_content=original_user_query,
-            past_messages=past_messages,
-            max_tokens=self.chat_token_limit - query_response_token_limit,  # TODO: count functions
-            fallback_to_default=True,
-        )
 
         chat_completion: ChatCompletion = await self.openai_chat_client.chat.completions.create(
             messages=query_messages,  # type: ignore
@@ -75,31 +74,32 @@ async def run(
             enable_text_search=text_search,
             filters=filters,
         )
-
-        sources_content = [f"[{(item.id)}]:{item.to_str_for_rag()}\n\n" for item in results]
+        
+        # Check if the url_filter is used to determine the context to send to the LLM
+        if any(f['column'] == 'url' and f['value'] != '' for f in filters):
+            sources_content = [f"[{(item.id)}]:{item.to_str_for_narrow_rag()}\n\n" for item in results] # all details
+        else:
+            sources_content = [f"[{(item.id)}]:{item.to_str_for_broad_rag()}\n\n" for item in results] # important details
+        
         content = "\n".join(sources_content)
 
-        # Generate a contextual and content specific answer using the search results and chat history
+        # Build messages for the final chat completion
+        answer_messages = copy.deepcopy(messages)
+        answer_messages.insert(0, {"role": "system", "content": self.answer_prompt_template})
+        answer_messages[-1]["content"].append({"type": "text", "text": "\n\nSources:\n" + content})
         response_token_limit = 1024
-        messages = build_messages(
-            model=self.chat_model,
-            system_prompt=overrides.get("prompt_template") or self.answer_prompt_template,
-            new_user_content=original_user_query + "\n\nSources:\n" + content,
-            past_messages=past_messages,
-            max_tokens=self.chat_token_limit - response_token_limit,
-            fallback_to_default=True,
-        )
 
         chat_completion_response = await self.openai_chat_client.chat.completions.create(
             # Azure OpenAI takes the deployment name as the model name
             model=self.chat_deployment if self.chat_deployment else self.chat_model,
-            messages=messages,
+            messages=answer_messages,
             temperature=overrides.get("temperature", 0.3),
             max_tokens=response_token_limit,
             n=1,
             stream=False,
         )
         chat_resp = chat_completion_response.model_dump()
+
         chat_resp["choices"][0]["context"] = {
             "data_points": {"text": sources_content},
             "thoughts": [