Update app.py

app.py

@@ -1,434 +1,97 @@
-import os
-import datetime
-import asyncio
+from src.llamaindex_palm import LlamaIndexPaLM, LlamaIndexPaLMText
 
-from typing import Any, List, Dict, Union
-from pydantic import Extra
-
-import wandb
-from wandb.sdk.data_types.trace_tree import Trace
-
-import pinecone
-import google.generativeai as genai
-
-from llama_index import (
-    ServiceContext,
-    PromptHelper,
-    VectorStoreIndex
-)
-from llama_index.vector_stores import PineconeVectorStore
-from llama_index.storage.storage_context import StorageContext
-from llama_index.node_parser import SimpleNodeParser
-from llama_index.text_splitter import TokenTextSplitter
-from llama_index.embeddings.base import BaseEmbedding
-from llama_index.llms import (
-    CustomLLM,
-    CompletionResponse,
-    CompletionResponseGen,
-    LLMMetadata,
-)
-from llama_index.llms.base import llm_completion_callback
-
-from llama_index.evaluation import SemanticSimilarityEvaluator
-from llama_index.embeddings import SimilarityMode
+import gradio as gr
 
+from typing import List
+import time
 import logging
-logging.basicConfig(format='%(asctime)s %(message)s', datefmt='%Y-%m-%d %I:%M:%S %p', level=logging.INFO)
-logger = logging.getLogger('llm')
-
-prompt_template = """
-[System]
-    You are in a role play of Gerard Lee.
-    Reply in no more than 7 complete sentences using content from [Context] only. Refer to [History] for seamless conversatation.
-
-[History]
-    {context_history}
 
-[Context]
-    {context_from_index}
-"""
+# Llama-Index LLM
+llm_backend = LlamaIndexPaLMText(model_kwargs={'temperature': 0.8})
+llm = LlamaIndexPaLM(model=llm_backend)
+llm.get_index_from_pinecone()
 
-class LlamaIndexPaLMEmbeddings(BaseEmbedding, extra=Extra.allow):
-    def __init__(
-        self,
-        model_name: str = 'models/embedding-gecko-001',
-        **kwargs: Any,
-    ) -> None:
-        super().__init__(**kwargs)
-        self._model_name = model_name
+# Gradio
+chat_history = []
 
-    @classmethod
-    def class_name(cls) -> str:
-        return 'PaLMEmbeddings'
+def clear_chat() -> None:
+    global chat_history 
+    chat_history = []
+    return None
 
-    def gen_embeddings(self, text: str) -> List[float]:
-        return genai.generate_embeddings(self._model_name, text)
+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 _get_query_embedding(self, query: str) -> List[float]:
-        embeddings = self.gen_embeddings(query)
-        return embeddings['embedding']
+def generate_text(prompt: str, llamaindex_llm: LlamaIndexPaLM):
+    global chat_history
 
-    def _get_text_embedding(self, text: str) -> List[float]:
-        embeddings = self.gen_embeddings(text)
-        return embeddings['embedding']
+    logger.info("Generating Message...")
+    logger.info(f"User Message:\n{prompt}\n")
 
-    def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
-        embeddings = [
-            self.gen_embeddings(text)['embedding'] for text in texts
-        ]
-        return embeddings
+    result = llamaindex_llm.generate_text(prompt, chat_history)
+    chat_history.append(prompt)
+    chat_history.append(result)
     
-    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)
-    
-class LlamaIndexPaLMText(CustomLLM, extra=Extra.allow):
-    def __init__(
-        self,
-        model_name: str = 'models/text-bison-001',
-        model_kwargs: dict = {},
-        context_window: int = 8196,
-        num_output: int = 1024,
-        **kwargs: Any,
-    ) -> None:
-        super().__init__(**kwargs)
-        self._model_name = model_name
-        self._model_kwargs = model_kwargs
-        self._context_window = context_window
-        self._num_output = num_output
-        
-    @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
-        )
-
-    def gen_texts(self, prompt):
-            logging.debug(f"prompt: {prompt}")
-            response = genai.generate_text(
-                model=self._model_name,
-                prompt=prompt,
-                safety_settings=[
-                    {
-                        'category': genai.types.HarmCategory.HARM_CATEGORY_UNSPECIFIED,
-                        'threshold': genai.types.HarmBlockThreshold.BLOCK_NONE,
-                    },
-                ],
-                **self._model_kwargs
+    logger.info(f"Replied Message:\n{result}\n")
+    return result
+
+if __name__ == "__main__":
+    logging.basicConfig(format='%(asctime)s %(message)s', datefmt='%Y-%m-%d %I:%M:%S %p', level=logging.INFO)
+    logger = logging.getLogger('app')
+
+    try:
+        with gr.Blocks() as app:
+            chatbot = gr.Chatbot(
+                bubble_full_width=False, 
+                container=True, 
+                show_share_button=False, 
+                avatar_images=[None, './asset/akag-g-only.png']
             )
-            logging.debug(f"response:\n{response}")
-            return response.candidates[0]['output']
-
-    @llm_completion_callback()
-    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
-        text = self.gen_texts(prompt)
-        return CompletionResponse(text=text)
-
-    @llm_completion_callback()
-    def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
-        raise NotImplementedError()
-    
-class LlamaIndexPaLM():
-    def __init__(
-        self,
-        emb_model: LlamaIndexPaLMEmbeddings = LlamaIndexPaLMEmbeddings(),
-        model: LlamaIndexPaLMText = LlamaIndexPaLMText(),
-        # prompt_template: str = prompt_template
-    ) -> None:
-        self.emb_model = emb_model
-        self.llm = model
-        self.prompt_template = prompt_template
-
-        # Google Generative AI
-        genai.configure(api_key=os.environ['PALM_API_KEY'])
-
-        # Pinecone
-        pinecone.init(
-            api_key=os.environ['PINECONE_API_KEY'],
-            environment=os.getenv('PINECONE_ENV')
-        )
-
-        # W&B
-        wandb.init(project=os.getenv('WANDB_PROJECT'))
-
-        # model metadata
-        CONTEXT_WINDOW = os.getenv('CONTEXT_WINDOW', 8196)
-        NUM_OUTPUT = os.getenv('NUM_OUTPUT', 1024)
-        TEXT_CHUNK_SIZE = os.getenv('TEXT_CHUNK_SIZE', 512)
-        TEXT_CHUNK_OVERLAP = os.getenv('TEXT_CHUNK_OVERLAP', 20)
-        TEXT_CHUNK_OVERLAP_RATIO = os.getenv('TEXT_CHUNK_OVERLAP_RATIO', 0.1)
-        TEXT_CHUNK_SIZE_LIMIT = os.getenv('TEXT_CHUNK_SIZE_LIMIT', None)
-
-        self.node_parser = SimpleNodeParser.from_defaults(
-            text_splitter=TokenTextSplitter(
-                chunk_size=TEXT_CHUNK_SIZE, chunk_overlap=TEXT_CHUNK_OVERLAP
+            msg = gr.Textbox(
+                show_label=False,
+                label="Type your message...", 
+                placeholder="Hi Gerard, can you introduce yourself?",
+                container=False, 
             )
-        )
-
-        self.prompt_helper = PromptHelper(
-            context_window=CONTEXT_WINDOW,
-            num_output=NUM_OUTPUT,
-            chunk_overlap_ratio=TEXT_CHUNK_OVERLAP_RATIO,
-            chunk_size_limit=TEXT_CHUNK_SIZE_LIMIT
-        )
-
-        self.service_context = ServiceContext.from_defaults(
-            llm=self.llm,
-            embed_model=self.emb_model,
-            node_parser=self.node_parser,
-            prompt_helper=self.prompt_helper,
-        )
-    
-        self.emd_evaluator = SemanticSimilarityEvaluator(
-            service_context=self.service_context,
-            similarity_mode=SimilarityMode.DEFAULT,
-            similarity_threshold=os.getenv('SIMILARITY_THRESHOLD', 0.7),
-        )
-
-    def get_index_from_pinecone(
-        self, 
-        index_name: str = os.getenv('PINECONE_INDEX'),
-        index_namespace: str = os.getenv('PINECONE_NAMESPACE')
-    ) -> None:
-        # Pinecone VectorStore
-        pinecone_index = pinecone.Index(index_name)
-        self.vector_store = PineconeVectorStore(pinecone_index=pinecone_index, add_sparse_vector=True, namespace=index_namespace)
-        self.pinecone_index = VectorStoreIndex.from_vector_store(self.vector_store, self.service_context)
-        self._index_name = index_name
-        self._index_namespace = index_namespace
-        return None
-    
-    async def retrieve_context(
-        self,
-        query: str
-    ) -> Dict[str, Union[str, int]]:
-        start_time = round(datetime.datetime.now().timestamp() * 1000)
-        response = await self.pinecone_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 evaluate_context(
-        self,
-        query: str,
-        returned_context: str
-    ) -> float:
-        result = await self.emd_evaluator.aevaluate(
-            response=returned_context,
-            reference=query,
-        )
-        return float(result.score)
-
-    def format_history_as_context(
-        self,
-        history: List[str],
-    ) -> str:
-        format_chat_history = "\n".join(list(filter(None, history)))
-        return format_chat_history
-
-    async def generate_text(
-        self,
-        query: str,
-        history: List[str],
-    ) -> str:
-        # get history
-        context_history = self.format_history_as_context(history=history)
-
-        # w&b trace start
-        start_time_ms = round(datetime.datetime.now().timestamp() * 1000)
-        root_span = Trace(
-            name="MetaAgent",
-            kind="agent",
-            start_time_ms=start_time_ms,
-            metadata={"user": "🤗 Space"},
-        )
-
-        # get retrieval context(s) from llama-index vectorstore index
-        # w&b trace retrieval & select agent
-        agent_span = Trace(
-            name="LlamaIndexAgent",
-            kind="agent",
-            start_time_ms=start_time_ms,
-        )
-        try:
-            # No history, single context retrieval without evaluation
-            if not history:
-                # w&b trace retrieval context
-                context_from_index_selected = self.retrieve_context(query)       
-                agent_end_time_ms = round(datetime.datetime.now().timestamp() * 1000)
-                retrieval_span = Trace(
-                    name="QueryRetrieval",
-                    kind="chain",
-                    status_code="success",
-                    metadata={
-                        "framework": "Llama-Index",
-                        "index_type": "VectorStoreIndex",
-                        "vector_store": "Pinecone",
-                        "vector_store_index": self._index_name,
-                        "vector_store_namespace": self._index_namespace,
-                        "model_name": self.llm._model_name,
-                        "custom_kwargs": self.llm._model_kwargs,
-                    },
-                    start_time_ms=start_time_ms,
-                    end_time_ms=agent_end_time_ms,
-                    inputs={"query": query},
-                    outputs={"response": context_from_index_selected},
-                )
-                agent_span.add_child(retrieval_span)
-            # Has history, multiple context retrieval with async, then evaluation to determine which context to choose
-            else:
-                extended_query = f"{history[-1]}\n{query}"
-                result_query_only, result_extended_query = await asyncio.gather(
-                    self.retrieve_context(query),
-                    self.retrieve_context(extended_query)
+            with gr.Row():
+                clear = gr.Button("Clear", scale=1)
+                send = gr.Button(
+                    value="", 
+                    variant="primary",
+                    icon="./asset/send-message.png", 
+                    scale=1
                 )
 
-                # w&b trace retrieval context query only
-                retrieval_query_span = Trace(
-                    name="QueryRetrieval",
-                    kind="chain",
-                    status_code="success",
-                    metadata={
-                        "framework": "Llama-Index",
-                        "index_type": "VectorStoreIndex",
-                        "vector_store": "Pinecone",
-                        "vector_store_index": self._index_name,
-                        "vector_store_namespace": self._index_namespace,
-                        "model_name": self.llm._model_name,
-                        "custom_kwargs": self.llm._model_kwargs,
-                    },
-                    start_time_ms=result_query_only.start,
-                    end_time_ms=result_query_only.end,
-                    inputs={"query": query},
-                    outputs={"response": result_query_only.result},
-                )
-                agent_span.add_child(retrieval_query_span)
+            def user(user_message, history):
+                return "", history + [[user_message, None]]
 
-                # w&b trace retrieval context extended query
-                retrieval_extended_query_span = Trace(
-                    name="ExtendedQueryRetrieval",
-                    kind="chain",
-                    status_code="success",
-                    metadata={
-                        "framework": "Llama-Index",
-                        "index_type": "VectorStoreIndex",
-                        "vector_store": "Pinecone",
-                        "vector_store_index": self._index_name,
-                        "vector_store_namespace": self._index_namespace,
-                        "model_name": self.llm._model_name,
-                        "custom_kwargs": self.llm._model_kwargs,
-                    },
-                    start_time_ms=result_extended_query.start,
-                    end_time_ms=result_extended_query.end,
-                    inputs={"query": extended_query},
-                    outputs={"response": result_extended_query.result},
-                )
-                agent_span.add_child(retrieval_extended_query_span)
-
-                # w&b trace select context
-                eval_start_time_ms = round(datetime.datetime.now().timestamp() * 1000)
-                eval_context_query_only, eval_context_extended_query = await asyncio.gather(
-                    self.evaluate_context(query, result_query_only.result),
-                    self.evaluate_context(extended_query, result_extended_query.result)
-                )
+            def bot(history):
+                bot_message = generate_text(history[-1][0], llm)
+                history[-1][1] = ""
+                for character in bot_message:
+                    history[-1][1] += character
+                    time.sleep(0.01)
+                    yield history
 
-                if eval_context_query_only > eval_context_extended_query:
-                    query_selected, context_from_index_selected = query, result_query_only.result 
-                else:
-                    query_selected, context_from_index_selected = extended_query, result_extended_query.result
-
-                agent_end_time_ms = round(datetime.datetime.now().timestamp() * 1000)
-                eval_span = Trace(
-                    name="EmbeddingsEvaluator",
-                    kind="tool",
-                    status_code="success",
-                    metadata={
-                        "framework": "Llama-Index",
-                        "evaluator": "SemanticSimilarityEvaluator",
-                        "similarity_mode": "DEFAULT",
-                        "similarity_threshold": 0.7,
-                        "similarity_results": {
-                            "eval_context_query_only": eval_context_query_only, 
-                            "eval_context_extended_query": eval_context_extended_query,
-                        },
-                        "model_name": self.emb_model._model_name,
-                    },
-                    start_time_ms=eval_start_time_ms,
-                    end_time_ms=agent_end_time_ms,
-                    inputs={"query": query_selected},
-                    outputs={"response": context_from_index_selected},
-                )
-                agent_span.add_child(eval_span)
-
-        except Exception as e:
-            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."
-            root_span.add_inputs_and_outputs(
-                inputs={"query": query}, outputs={"result": result, "exception": e}
+            msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+                bot, chatbot, chatbot
             )
-            root_span._span.status_code="fail"
-            root_span._span.end_time_ms = llm_end_time_ms
-            root_span.log(name="llm_app_trace")
-            return result
-        
-        logger.info(f"Context from Llama-Index:\n{context_from_index_selected}\n")
-
-        agent_span.add_inputs_and_outputs(
-            inputs={"query": query}, outputs={"result": context_from_index_selected}
-        )
-        agent_span._span.status_code="success"
-        agent_span._span.end_time_ms = agent_end_time_ms
-        root_span.add_child(agent_span)
-
-        # 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:
-            response = genai.generate_text(
-                prompt=prompt_with_context,
-                safety_settings=[
-                    {
-                        'category': genai.types.HarmCategory.HARM_CATEGORY_UNSPECIFIED,
-                        'threshold': genai.types.HarmBlockThreshold.BLOCK_NONE,
-                    },
-                ],
-                temperature=0.9,
+            send.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+                bot, chatbot, chatbot
             )
-            result = response.result
-            success_flag = "success"
-            if result is None:
-                result = "Seems something went wrong. Please try again later."
-                logger.error(f"Result with 'None' received\n")
-                success_flag = "fail"
-
-        except Exception as e:
-            result = "Seems something went wrong. Please try again later."
-            logger.error(f"Exception {e} occured\n")
-            success_flag = "fail"
-
-        # w&b trace llm
-        llm_end_time_ms = round(datetime.datetime.now().timestamp() * 1000)
-        llm_span = Trace(
-            name="LLM",
-            kind="llm",
-            status_code=success_flag,
-            start_time_ms=agent_end_time_ms,
-            end_time_ms=llm_end_time_ms,
-            inputs={"input": prompt_with_context},
-            outputs={"result": result},
-        )
-        root_span.add_child(llm_span)
-
-        # w&b finalize trace
-        root_span.add_inputs_and_outputs(
-            inputs={"query": query}, outputs={"result": result}
-        )
-        root_span._span.end_time_ms = llm_end_time_ms
-        root_span.log(name="llm_app_trace")
-
-        return result
+            clear.click(clear_chat, None, chatbot, queue=False)
+
+            gr.HTML("""
+                <p><center><i>Disclaimer: This is a RAG app for demostration purpose. LLM hallucination might occur.</i></center></p>
+                <p><center>Hosted on 🤗 Spaces. Powered by Google PaLM 🌴</center></p>
+            """)
+            
+        app.queue()
+        app.launch()
+    except Exception as e:
+        logger.exception(e)