Update requirements.txt to include llama-cpp-python dependency; change default port in launch.json from 8000 to 8080; add VSCode settings for Python type checking; modify welcome message in main.py; enhance configuration in config.py with new model and file name; implement Message and ChatResponse models for structured messaging; refactor chat_request and chat_service to utilize new message structure; streamline chat response handling; and update client.py for improved OpenAI API integration.

.vscode/launch.json

@@ -24,7 +24,7 @@
                 "src.main:app",
                 "--reload",
                 "--port",
-                "8000",
+                "8080",
                 "--host",
                 "0.0.0.0",
             ]

.vscode/settings.json

@@ -0,0 +1,4 @@
+{
+  "python.analysis.typeCheckingMode": "basic",
+  "python.analysis.autoImportCompletions": true
+}

requirements.txt

@@ -21,4 +21,5 @@ langchain_chroma>=0.2.2
 chromadb>=0.6.3
 sentence_transformers>=4.1.0
 langchain_huggingface>=0.1.2
-huggingface_hub[hf_xet]
+huggingface_hub[hf_xet]
+llama-cpp-python==0.3.8

src/constants/config.py

@@ -7,9 +7,12 @@ TORCH_DEVICE = (
     else "cpu"
 )
 IMAGE_MODEL_ID_OR_LINK = "stable-diffusion-v1-5/stable-diffusion-v1-5"
+MODEL_NAME = "facebook/opt-125m"
 CACHE_DIR = "/tmp/cache"
 DATA_DIR = "/tmp/data"
 EMBEDDING_MODEL = "intfloat/multilingual-e5-large-instruct"
 UPLOAD_DIR = "/tmp/uploads"
 OUTPUT_DIR = "/tmp/outputs"
+FILE_NAME = "super-lite-model.gguf"
+# FILE_NAME = "llama_3.1_8b_instruct_q4_k_m.gguf"
 # EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

src/main.py

@@ -47,7 +47,7 @@ app.include_router(process_file_routes.router, prefix="/api/v1")
 app.include_router(vector_store_routes.router, prefix="/api/v1")
 @app.get("/")
 def read_root():
-    return {"message": "Welcome my API"}
+    return {"message": "Welcome to my API"}
 
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 app.mount(OUTPUT_DIR, StaticFiles(directory=OUTPUT_DIR), name="outputs")

src/models/others/message.py

@@ -0,0 +1,22 @@
+from enum import Enum
+from typing import Optional
+
+from pydantic import BaseModel
+
+
+class Role(str, Enum):
+    assistant = "assistant"
+    user = "user"
+    system = "system"
+    tool = "tool"
+
+
+class Message(BaseModel):
+    role: Role
+    content: Optional[str] = None
+
+    def to_map(self):
+        return {
+            "role": self.role.value,
+            "content": self.content,
+        }

src/models/requests/chat_request.py

@@ -1,18 +1,27 @@
+from typing import List, Optional
 from pydantic import BaseModel
 
+from constants.config import MODEL_NAME
+from models.others.message import Role, Message
+
+
 class ChatRequest(BaseModel):
-    prompt: list[dict]
+    messages: List[Message]
+    # temperature: Optional[float] = 0.7
+    # max_tokens: Optional[int] = -1
     has_file: bool = False
     chat_session_id: str | None = None
-    
+
     model_config = {
         "json_schema_extra": {
             "examples": [
                 {
-                    "prompt": [{"role": "user", "content": "Hello, how are you?"}],
                     "has_file": False,
-                    "chat_session_id": "123"
+                    "chat_session_id": "123",
+                    "messages": [{"role": Role.user, "content": "hello"}],
+                    # "temperature":0.7,
+                    # "max_tokens": -1,
                 }
             ]
         }
-    }
+    }

src/models/responses/chat_response.py

@@ -0,0 +1,70 @@
+from typing import Any, List, Optional
+from pydantic import BaseModel
+from models.others.message import Message, Role
+
+# class Usage(BaseModel):
+#     prompt_token: int
+#     completion_token: int
+#     total_tokens: int
+
+
+class Choice(BaseModel):
+    # index: int
+    # logprobs: Any
+    # finish_reason: Optional[str]
+    message: Optional[Message] = None
+    delta: Optional[Message] = None
+
+
+class ChatResponse(BaseModel):
+    id: Optional[str] = None
+    # object: Optional[str] = None
+    # created: Optional[int] = None
+    # model: Optional[str] = None
+    # system_fingerprint: Optional[str] = None
+    # usage: Optional[Usage] = None
+    choices: Optional[List[Choice]] = None
+
+    @classmethod
+    def from_stream_chunk(cls, chunk: dict, last_role: Optional[Role] = None):
+        choices = []
+        updated_role = last_role  # Default to last role
+
+        for choice in chunk.get("choices", []):
+            delta_data = choice.get("delta", {})
+
+            # Skip chunks that contain neither content nor role
+            if not delta_data.get("content") and not delta_data.get("role"):
+                continue
+
+            # Determine role
+            if "role" in delta_data and delta_data["role"] is not None:
+                try:
+                    updated_role = Role(delta_data["role"])
+                except ValueError:
+                    # Skip or log invalid role values
+                    continue
+
+            if not updated_role:
+                # Still no role available, skip
+                continue
+
+            message = Message(
+                role=updated_role,
+                content=delta_data.get("content"),
+            )
+
+            choices.append(
+                Choice(
+                    message=message,
+                    delta=message,
+                )
+            )
+
+        return (
+            cls(
+                id=chunk.get("id"),
+                choices=choices,
+            ),
+            updated_role,
+        )

src/routes/chat_routes.py

@@ -9,7 +9,6 @@ from models.requests.chat_request import ChatRequest
 from models.responses.base_exception_response import BaseExceptionResponse
 from models.responses.base_response import BaseResponse
 from services import chat_service
-from services.process_file_service import get_file_content
 
 router = APIRouter(tags=["Chat"])
 
@@ -63,7 +62,7 @@ async def chat(request: ChatRequest):
 
     try:
         response = chat_service.chat_generate(request=request)
-        return BaseResponse(data=json.loads(response.model_dump_json()))
+        return BaseResponse(data=response)
     except Exception as e:
         raise BaseExceptionResponse(message=str(e))
     

src/routes/vector_store_routes.py

@@ -1,13 +1,7 @@
-import json
-import uuid
 from fastapi import APIRouter
-from models.requests.chat_request import ChatRequest
 from models.responses.base_exception_response import BaseExceptionResponse
 from models.responses.base_response import BaseResponse
 from services import  vector_store_service
-from utils.client import openai_client
-import os
-from chromadb import PersistentClient
 
 
 router = APIRouter(tags=["Vector Store"])

src/services/chat_service.py

@@ -1,30 +1,34 @@
 from constants import system_prompts
 from models.requests.chat_request import ChatRequest
 from services import vector_store_service
+from utils.client import create, create_stream
 from utils.timing import measure_time
 from utils.tools import tools_helper, tools_define
+from models.others.message import Message, Role
 
-def build_context_prompt(request: ChatRequest) -> list:
+
+def build_context_prompt(request: ChatRequest) -> list[Message]:
     """Build system prompt with context if file is provided."""
-    messages = [{"role": "system", "content": system_prompts.system_prompt}]
-    
-    if not request.has_file or not vector_store_service.check_if_collection_exists(request.chat_session_id):
-        return messages 
-    
+    messages = [Message(role=Role.system, content=system_prompts.system_prompt)]
+
+    if not request.has_file or not vector_store_service.check_if_collection_exists(
+        request.chat_session_id
+    ):
+        return messages
 
     with measure_time("Get data from vector store"):
         vectorstore = vector_store_service.get_vector_store(request.chat_session_id)
-        query = request.prompt[-1]["content"]
-        results = vectorstore.similarity_search(query=query, k=10)
+        query = request.messages[-1].content
+        results = vectorstore.similarity_search(query=query or "", k=10)
 
     if not results:
         return messages
 
     with measure_time("Building context prompt"):
-        context = ''
+        context = ""
     for document in results:
         # print(f"Document:{document.page_content[:50]}, score:{score}\n\n")
-        source = document.metadata.get('file_id', 'Unknown File')
+        source = document.metadata.get("file_id", "Unknown File")
         context += f"Context from file: {source}\n\n{document.page_content}\n\n"
 
     embedded_prompt = (
@@ -35,56 +39,62 @@ def build_context_prompt(request: ChatRequest) -> list:
         f"CONTEXT: {context}\nQUESTION: {query}"
     )
 
-    messages.append({"role": "system", "content": embedded_prompt})
+    messages.append(Message(role=Role.system, content=embedded_prompt))
     return messages
 
-def chat_generate_stream(request: ChatRequest,):
+
+def chat_generate_stream(
+    request: ChatRequest,
+):
     """Streaming chat generation."""
     messages = build_context_prompt(request)
-    messages.extend(request.prompt)
+    messages.extend(request.messages)
 
-    stream = openai_client.chat.completions.create(
-        messages=messages,
-        model='my-model',
-        stream=True,
-        tools=tools_define.tools
-    )
+    # stream = openai_client.chat.completions.create(
+    #     messages=messages,
+    #     model='my-model',
+    #     stream=True,
+    #     tools=tools_define.tools
+    # )
+
+    stream = create_stream(messages)
 
     final_tool_calls = {}
 
     for chunk in stream:
-        delta = chunk.choices[0].delta
-        if getattr(delta, 'tool_calls', None):
-            final_tool_calls = tools_helper.final_tool_calls_handler(final_tool_calls, delta.tool_calls)
-        yield chunk
+        if chunk.choices and len(chunk.choices) > 0:
+            delta = chunk.choices[0].delta
+            if getattr(delta, "tool_calls", None):
+                final_tool_calls = tools_helper.final_tool_calls_handler(
+                    final_tool_calls, delta.tool_calls
+                )
+            yield chunk
 
     if not final_tool_calls:
         return
 
-    tool_call_message = tools_helper.process_tool_calls(final_tool_calls)
+    tool_call_result = tools_helper.process_tool_calls(final_tool_calls)
+    tool_call_message = Message(
+        role=Role.tool, content=tool_call_result.get("content", "")
+    )
     messages.append(tool_call_message)
 
-    new_stream = openai_client.chat.completions.create(
-        messages=messages,
-        model='my-model',
-        stream=True
-    )
+    new_stream = create_stream(messages)
 
     for chunk in new_stream:
         yield chunk
 
 
-def chat_generate(request: ChatRequest,):
+def chat_generate(request: ChatRequest):
     """Non-streaming (batched) chat generation."""
     messages = build_context_prompt(request)
-    messages.extend(request.prompt)
+    messages.extend(request.messages)
 
     with measure_time("Non-streaming chat generation"):
-        response = openai_client.chat.completions.create(
-            messages=messages,
-            model='my-model',
-            tools=tools_define.tools
-        )
-        return response
-        
-
+        # response = openai_client.chat.completions.create(
+        #     messages=messages,
+        #     model='my-model',
+        #     tools=tools_define.tools
+        # )
+        output = create(messages=messages)
+        return output

src/utils/client.py

@@ -1,8 +1,50 @@
+import os
+from typing import Generator, List
 import openai
+import torch
+
+from constants.config import FILE_NAME
+from models.others.message import Message
+from models.requests.chat_request import ChatRequest
+from models.responses.chat_response import ChatResponse
+from utils.tools import tools_define
 
 # Initialize OpenAI API client
 openai_client = openai.OpenAI(
-    base_url="http://localhost:1234/v1",
+    base_url="http://localhost:8000/v1",
     api_key="none",
 )
 
+from llama_cpp import ChatCompletionTool, Llama
+
+# Determine number of CPU threads based on device
+if torch.cuda.is_available() or (
+    hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+):
+    n_threads = 4  # Fewer threads if using GPU/MPS, adjust as needed
+    n_gpu_layers = 20
+else:
+    n_threads = os.cpu_count() or 4
+    n_gpu_layers = 0
+
+# Khởi tạo mô hình từ GGUF
+llm = Llama(
+    model_path=FILE_NAME,
+    n_threads=n_threads,
+    n_gpu_layers=n_gpu_layers,
+    n_ctx=4096,
+)
+
+def create(messages: List[Message]):
+    prompt = [message.to_map() for message in messages]
+    output = llm.create_chat_completion(prompt)  # type: ignore
+    return output
+
+def create_stream(messages: List[Message]) -> Generator[ChatResponse, None, None]:
+    prompt = [message.to_map() for message in messages]
+    output = llm.create_chat_completion(prompt, stream=True, tools=tools_define.tools)  # type: ignore
+    last_role = None
+    for chunk in output:
+        response, last_role = ChatResponse.from_stream_chunk(chunk, last_role)  # type: ignore
+        if response.choices:
+            yield response

src/utils/image_pipeline.py

@@ -1,32 +1,32 @@
-import torch
-from diffusers import StableDiffusionPipeline
-from constants.config import IMAGE_MODEL_ID_OR_LINK, TORCH_DEVICE
+# import torch
+# from diffusers import StableDiffusionPipeline
+# from constants.config import IMAGE_MODEL_ID_OR_LINK, TORCH_DEVICE
 
-torch.backends.cuda.matmul.allow_tf32 = True  # Enable TF32 for performance on CUDA
+# torch.backends.cuda.matmul.allow_tf32 = True  # Enable TF32 for performance on CUDA
 
-_pipeline = None
+# _pipeline = None
 
-def get_pipeline() -> StableDiffusionPipeline:
-    global _pipeline
-    if _pipeline is None:
-        try:
-            _pipeline = StableDiffusionPipeline.from_pretrained(
-                IMAGE_MODEL_ID_OR_LINK,
-                torch_dtype=torch.bfloat16,
-                variant="fp16",
-                # safety_checker=True,
-                use_safetensors=True,
-            )
-            # _pipeline = StableDiffusionPipeline.from_single_file(
-            #     IMAGE_MODEL_ID_OR_LINK,
-            #     torch_dtype=torch.bfloat16,
-            #     variant="fp16",
-            #     # safety_checker=True,
-            #     use_safetensors=True,
-            # )
-            _pipeline.to(TORCH_DEVICE)
-        except Exception as e:
-            raise RuntimeError(f"Failed to load the model: {e}")
-    return _pipeline
+# def get_pipeline() -> StableDiffusionPipeline:
+#     global _pipeline
+#     if _pipeline is None:
+#         try:
+#             _pipeline = StableDiffusionPipeline.from_pretrained(
+#                 IMAGE_MODEL_ID_OR_LINK,
+#                 torch_dtype=torch.bfloat16,
+#                 variant="fp16",
+#                 # safety_checker=True,
+#                 use_safetensors=True,
+#             )
+#             # _pipeline = StableDiffusionPipeline.from_single_file(
+#             #     IMAGE_MODEL_ID_OR_LINK,
+#             #     torch_dtype=torch.bfloat16,
+#             #     variant="fp16",
+#             #     # safety_checker=True,
+#             #     use_safetensors=True,
+#             # )
+#             _pipeline.to(TORCH_DEVICE)
+#         except Exception as e:
+#             raise RuntimeError(f"Failed to load the model: {e}")
+#     return _pipeline
 
-pipeline = get_pipeline()
+# pipeline = get_pipeline()