Upload 9 files

app.py

@@ -0,0 +1,114 @@
+from src.config import setup_logging, Config
+from src.embeddings_model import GEmbeddings
+from src.text_generation_model import GLLM
+from src.pinecone_index import PineconeIndex
+from src.llamaindex_backend import GLlamaIndex
+
+import gradio as gr
+import google.generativeai as genai
+from llama_index.core import Settings
+
+from typing import List
+import time
+
+import dotenv
+dotenv.load_dotenv(".env")
+
+logger = setup_logging()
+
+# Google Generative AI
+genai.configure(api_key=Config.GAI_API_KEY)
+
+# Llama-Index LLM
+embed_model = GEmbeddings(model_name=Config.EMB_MODEL_NAME)
+llm = GLLM(model_name=Config.TEXT_MODEL_NAME, system_instruction=None)
+
+Settings.embed_model = embed_model
+Settings.llm = llm
+
+index = PineconeIndex(api_key=Config.PINECONE_API_KEY, index_name=Config.PC_INDEX_NAME, index_namespace=Config.PC_INDEX_NAMESPACE)
+backend = GLlamaIndex(logger, embed_model, llm, index, Config.SIMILARITY_THRESHOLD)
+
+# Gradio
+chat_history = []
+
+def clear_chat() -> None:
+    global chat_history 
+    chat_history = []
+    return None
+
+def get_chat_history(chat_history: List[str]) -> str:
+    ind = 0
+    formatted_chat_history = ""
+    for message in chat_history:
+        formatted_chat_history += f"User: \n{message}\n" if ind % 2 == 0 else f"Bot: \n{message}\n"
+        ind += 1
+    return formatted_chat_history
+
+def generate_text(prompt: str, backend: GLlamaIndex):
+    global chat_history
+
+    logger.info("Generating Message...")
+    logger.info(f"User Message:\n{prompt}\n")
+
+    result = backend.generate_text(prompt, chat_history)
+    chat_history.append(prompt)
+    chat_history.append(result)
+    
+    logger.info(f"Replied Message:\n{result}\n")
+    return result
+
+if __name__ == "__main__":
+    try:
+        with gr.Blocks(css=".input textarea {font-size: 16px !important}") as app:
+            chatbot = gr.Chatbot(
+                bubble_full_width=False, 
+                container=True, 
+                show_share_button=False, 
+                avatar_images=[None, './asset/akag-g-only.png']
+            )
+            msg = gr.Textbox(
+                show_label=False,
+                label="Type your message...", 
+                placeholder="Hi Gerard, can you introduce yourself?",
+                container=False, 
+                elem_classes="input"
+            )
+            with gr.Row():
+                clear = gr.Button("Clear", scale=1)
+                send = gr.Button(
+                    value="", 
+                    variant="primary",
+                    icon="./asset/send-message.png", 
+                    scale=1
+                )
+
+            def user(user_message, history):
+                return "", history + [[user_message, None]]
+
+            def bot(history):
+                bot_message = generate_text(history[-1][0], backend)
+                history[-1][1] = ""
+                for character in bot_message:
+                    history[-1][1] += character
+                    time.sleep(0.01)
+                    yield history
+
+            msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+                bot, chatbot, chatbot
+            )
+            send.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+                bot, chatbot, chatbot
+            )
+            clear.click(clear_chat, None, chatbot, queue=False)
+
+            gr.HTML("""
+                <p><center><i>Disclaimer: This RAG app is for demostration only. Hallucination might occur.</i></center></p>
+                <p><center>Hosted on 🤗 Spaces | Built with Google Gemini & 🦙 LlamaIndex | Last updated 2025</center></p>
+            """)
+            
+        app.queue()
+        app.launch()
+        
+    except Exception as e:
+        logger.exception(e)

requirements.txt

@@ -0,0 +1,9 @@
+google-generativeai==0.8.3
+llama-index==0.12.12
+llama-index-vector-stores-pinecone==0.4.2
+transformers==4.48.0
+pinecone-client==5.0.1
+wandb==0.19.2
+# transformers==4.30.2
+# llama-index==0.8.48
+# wandb==0.15.12

src/config.py

@@ -0,0 +1,30 @@
+import os
+import logging
+
+def setup_logging():
+    logging.basicConfig(
+        format='%(asctime)s %(message)s',
+        datefmt='%Y-%m-%d %I:%M:%S %p',
+        level=logging.INFO
+    )
+    return logging.getLogger('backend')
+
+class Config:
+    # External services
+    GAI_API_KEY = os.environ['GAI_API_KEY']
+    PINECONE_API_KEY = os.environ['PINECONE_API_KEY']
+    WANDB_API_KEY = os.environ['WANDB_API_KEY']
+    WANDB_PROJECT = os.environ['WANDB_PROJECT']
+    
+    # Model settings
+    TEXT_MODEL_NAME = os.getenv('TEXT_MODEL_NAME', 'gemini-1.5-flash')
+    EMB_MODEL_NAME = os.getenv('EMB_MODEL_NAME', 'models/text-embedding-004')
+    PC_INDEX_NAME = os.getenv('PC_INDEX_NAME', 'main-index')
+    PC_INDEX_NAMESPACE = os.getenv('PC_INDEX_NAMESPACE', 'main')
+    CONTEXT_WINDOW = int(os.getenv('CONTEXT_WINDOW', 32768))
+    NUM_OUTPUT = int(os.getenv('NUM_OUTPUT', 4098))
+    TEXT_CHUNK_SIZE = int(os.getenv('TEXT_CHUNK_SIZE', 2048))
+    TEXT_CHUNK_OVERLAP = int(os.getenv('TEXT_CHUNK_OVERLAP', 200))
+    TEXT_CHUNK_OVERLAP_RATIO = float(os.getenv('TEXT_CHUNK_OVERLAP_RATIO', 0.1))
+    TEXT_CHUNK_SIZE_LIMIT = os.getenv('TEXT_CHUNK_SIZE_LIMIT', None)
+    SIMILARITY_THRESHOLD = float(os.getenv('SIMILARITY_THRESHOLD', 0.7))

src/embeddings_model.py

@@ -0,0 +1,35 @@
+from typing import List, Any
+import google.generativeai as genai
+from llama_index.core.embeddings import BaseEmbedding
+
+class GEmbeddings(BaseEmbedding):
+    def __init__(
+        self,
+        model_name: str = 'models/text-embedding-004',
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self._model_name = model_name
+
+    def gai_embed_content(self, text: str) -> List[float]:
+        return genai.embed_content(model=self._model_name, content=text)
+
+    def _get_query_embedding(self, query: str) -> List[float]:
+        embeddings = self.gai_embed_content(query)
+        return embeddings['embedding']
+
+    def _get_text_embedding(self, text: str) -> List[float]:
+        embeddings = self.gai_embed_content(text)
+        return embeddings['embedding']
+
+    def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
+        embeddings = [
+            self.gai_embed_content(text)['embedding'] for text in texts
+        ]
+        return embeddings
+
+    async def _aget_query_embedding(self, query: str) -> List[float]:
+        return self._get_query_embedding(query)
+
+    async def _aget_text_embedding(self, text: str) -> List[float]:
+        return self._get_text_embedding(text)

src/llamaindex_backend.py

@@ -0,0 +1,145 @@
+from src.embeddings_model import GEmbeddings
+from src.text_generation_model import GLLM
+from src.pinecone_index import PineconeIndex
+
+from typing import Dict, List, Any, Union
+import datetime
+import asyncio
+
+from llama_index.core.evaluation import SemanticSimilarityEvaluator
+from llama_index.core.base.embeddings.base import SimilarityMode
+
+prompt_template = """
+<system>
+You are in a role play of Gerard Lee. Gerard is a data enthusiast and humble about his success.
+Reply as faifhfully as possible and in no more than 5 complete sentences unless <user query> requests to elaborate in details. Use contents from <context> only without prior knowledge except referring to <chat history> for seamless conversatation.
+</system>
+
+<chat history>
+{context_history}
+<chat history>
+
+<context>
+{context_from_index}
+</context>
+
+<user query>
+{user_query}
+</user query>
+"""
+
+class GLlamaIndex():
+    def __init__(
+        self,
+        logger,
+        emb_model: GEmbeddings,
+        text_model: GLLM,
+        index: PineconeIndex,
+        similarity_threshold: float
+    ) -> None:
+        self.logger = logger
+        self.emb_model = emb_model
+        self.llm = text_model
+        self.index = index
+        self.evaluator = self._set_evaluator(similarity_threshold)
+        self.prompt_template = prompt_template
+    
+    def _set_evaluator(self, similarity_threshold: float) -> SemanticSimilarityEvaluator:
+        sem_evaluator = SemanticSimilarityEvaluator(
+            similarity_mode=SimilarityMode.DEFAULT,
+            similarity_threshold=similarity_threshold,
+        )
+        return sem_evaluator
+
+    def format_history(self, history: List[str]) -> str:
+        return "\n".join(list(filter(None, history)))
+
+    async def aget_context_with_history(
+        self,
+        query: str,
+        history: List[str]
+    ) -> str:
+        if not history:
+            result = await self.index.retrieve_context(query)
+            return result["result"]
+            
+        extended_query = f"[History]\n{history[-1]}\n[New Query]\n{query}"
+        results = await self.index.aretrieve_context_multi(
+            [query, extended_query]
+        )
+        print(results)
+        eval_results = await self.aevaluate_context_multi(
+            [query, extended_query],
+            [r["result"] for r in results]
+        )
+        print(eval_results)
+        return results[0]["result"] if eval_results[0].score > eval_results[1].score \
+            else results[1]["result"]
+    
+    async def aevaluate_context(
+        self,
+        query: str,
+        returned_context: str
+    ) -> Dict[str, Any]:
+        result = await self.evaluator.aevaluate(
+            response=returned_context,
+            reference=query,
+        )
+        return result
+
+    async def aevaluate_context_multi(
+        self,
+        query_list: List[str],
+        returned_context_list: List[str]
+    ) -> List[Dict]:
+        result = await asyncio.gather(*(self.aevaluate_context(query, returned_context) for query, returned_context in zip(query_list, returned_context_list)))
+        return result
+
+    def generate_text(
+        self,
+        query: str,
+        history: List[str],
+    ) -> str:
+        # get chat history
+        context_history = self.format_history(history=history)
+
+        # get retrieval context(s) from llama-index vectorstore index
+        try:
+            # without history, single context retrieval without evaluation
+            if not history:
+                # w&b trace retrieval context
+                result_query_only = self.index.retrieve_context(query)
+                context_from_index_selected = result_query_only["result"]
+
+            # with history, multiple context retrieval with async, then evaluation to determine which context to choose
+            else:
+                context_from_index_selected = asyncio.run(self.aget_context_with_history(query=query, history=history))
+
+        except Exception as e:
+            self.logger.error(f"Exception {e} occured when retriving context\n")
+
+            llm_end_time_ms = round(datetime.datetime.now().timestamp() * 1000)
+            result = "Something went wrong. Please try again later."
+            return result
+        
+        self.logger.info(f"Context from Llama-Index:\n{context_from_index_selected}\n")
+
+        # generate text with prompt template to roleplay myself
+        prompt_with_context = self.prompt_template.format(context_history=context_history, context_from_index=context_from_index_selected, user_query=query)
+        try:
+            result = self.llm.gai_generate_content(
+                prompt=prompt_with_context,
+                temperature=0.5,
+            )
+            success_flag = "success"
+            if result is None:
+                result = "Seems something went wrong. Please try again later."
+                self.logger.error(f"Result with 'None' received\n")
+                success_flag = "fail"
+
+        except Exception as e:
+            result = "Seems something went wrong. Please try again later."
+            self.logger.error(f"Exception {e} occured\n")
+            success_flag = "fail"
+
+        return result

src/pinecone_index.py

@@ -0,0 +1,42 @@
+from typing import List, Dict, Union
+import datetime
+import asyncio
+from pinecone import Pinecone
+from llama_index.core import VectorStoreIndex, StorageContext
+from llama_index.vector_stores.pinecone import PineconeVectorStore
+
+class PineconeIndex:
+    def __init__(self, api_key: str, index_name: str, index_namespace: str):
+        self._index_name = index_name
+        self._index_namespace = index_namespace
+        self._pc = Pinecone(api_key=api_key)
+        self.pc_index = self._set_index(index_name, index_namespace)
+    
+    def _set_index(self, index_name: str, index_namespace: str) -> VectorStoreIndex:
+        vector_store = PineconeVectorStore(
+            pinecone_index=self._pc.Index(index_name),
+            add_sparse_vector=True, 
+            namespace=index_namespace
+        )
+        storage_context = StorageContext.from_defaults(vector_store=vector_store)
+        pc_index = VectorStoreIndex.from_vector_store(vector_store=vector_store, storage_context=storage_context)
+        return pc_index
+
+    def retrieve_context(self, query: str) -> Dict[str, Union[str, int]]:
+        start_time = round(datetime.datetime.now().timestamp() * 1000)
+        response = self.pc_index.as_query_engine(similarity_top_k=3).query(query)
+        end_time = round(datetime.datetime.now().timestamp() * 1000)
+        return {"result": response.response, "start": start_time, "end": end_time}
+
+    async def aretrieve_context(self, query: str) -> Dict[str, Union[str, int]]:
+        start_time = round(datetime.datetime.now().timestamp() * 1000)
+        response = await self.pc_index.as_query_engine(
+            similarity_top_k=3,
+            use_async=True
+        ).aquery(query)
+        end_time = round(datetime.datetime.now().timestamp() * 1000)
+        return {"result": response.response, "start": start_time, "end": end_time}
+
+    async def aretrieve_context_multi(self, query_list: List[str]) -> List[Dict]:
+        result = await asyncio.gather(*(self.aretrieve_context(query) for query in query_list))
+        return result

src/text_generation_model.py

@@ -0,0 +1,61 @@
+from typing import Any
+import google.generativeai as genai
+from google.generativeai.types import HarmCategory, HarmBlockThreshold
+from llama_index.core.llms import (
+    CustomLLM,
+    CompletionResponse,
+    CompletionResponseGen,
+    LLMMetadata,
+)
+from llama_index.core.llms.callbacks import llm_completion_callback
+
+class GLLM(CustomLLM):
+    def __init__(
+        self,
+        context_window: int = 32768,
+        num_output: int = 4098,
+        model_name: str = "gemini-1.5-flash",
+        system_instruction: str = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self._context_window = context_window
+        self._num_output = num_output
+        self._model_name = model_name
+        self._model = genai.GenerativeModel(model_name, system_instruction=system_instruction)
+
+    def gai_generate_content(self, prompt: str, temperature:float =0.5) -> str:
+        return self._model.generate_content(
+            prompt,
+            generation_config = genai.GenerationConfig(
+                temperature=temperature,
+            ),
+            safety_settings={
+                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
+            }
+        ).text
+
+    @property
+    def metadata(self) -> LLMMetadata:
+        """Get LLM metadata."""
+        return LLMMetadata(
+            context_window=self._context_window,
+            num_output=self._num_output,
+            model_name=self._model_name,
+        )
+
+    @llm_completion_callback()
+    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+        text = self.gai_generate_content(prompt)
+        return CompletionResponse(text=text)
+
+    @llm_completion_callback()
+    def stream_complete(
+        self, prompt: str, **kwargs: Any
+    ) -> CompletionResponseGen:
+        text = self.gai_generate_content(prompt)
+        response = ""
+        for token in text:
+            response += token
+            yield CompletionResponse(text=response, delta=token)