working demo with llm

.gitignore

@@ -83,4 +83,3 @@ dmypy.json
 .pyre/
 .pytype/
 cython_debug/
-temp.pdf

.projen/deps.json

@@ -14,7 +14,7 @@
     },
     {
       "name": "projen",
-      "version": "0.71.71",
+      "version": "0.71.72",
       "type": "devenv"
     },
     {

requirements-dev.txt

@@ -2,7 +2,7 @@
 black
 flake8
 mypy
-projen==0.71.71
+projen==0.71.72
 pylint
 pytest
 streamlit

requirements.txt

@@ -9,4 +9,3 @@ pydantic[dotenv]
 streamlit_chat
 tiktoken
 youtube-transcript-api
-loguru

taai_demo/plugin/{chat_response.py → llm.py}

@@ -4,16 +4,42 @@ from pydantic import BaseModel, PrivateAttr, validator
 import openai
 import tiktoken
 from loguru import logger
-from src.runtime_settings import Settings
+from taai_demo.settings import Settings
 
 
-SEARCH_SYSTEM_PROMPT = "You a search engine. You will be given information and a question and you will need to answer the question using the information. If the information provided is not enough or if I don't provide any information, you must respond: 'No Results Found'. Remember if the information I provide is not relevant or if I don't provide any information to you, you must respond: 'No Results Found' and nothing else. You must only only respond with 'No Results Found'."
-SEARCH_CONTEXT_PROMPT = """I found some information to help you. Please read it and use it to help you respond. If the information provided is not enough, you must respond: 'No Results Found'
-
-Information: {context}
-Query: {question}
+SEARCH_SYSTEM_PROMPT = """
+You a search engine. You will be provided a search query and you should then find relevant information to answer the query for the user.
+The information you find may come from a pdf or from a youtube video. If from a Youtube video, you will know this as the
+a metadata field will be returned with the youtube video id, start_time and end_time. If any of the information you use in your response to the user query
+comes from the youtube video, you must return a bulleted list of youtube video urls (format: https://www.youtube.com/watch?v=video_id&t=95sstart_time -> where start_time is formatted like 300s).
+so that the user can find the information. This will ensure the user isn't plagerizing and is using the information to learn.
+Finally, you MUST use markdown format when you respond to the user to help them format their response. Any links
+should be clickable using markdown format. Remember, if you can't find something, please respond with {{No information found}}.
+"""
+SEARCH_CONTEXT_PROMPT = """Thought: I should search for information related to the users query.
+Action: Search database for information relevant to the user query.
+Information returned from the database: {context}
+Here's my final response to the user query: {query}
+FORMAT for response:
+{{Response}}
+{{citations from the metadata field returned from database}} -> if the citation is a youtube video (denoted by the metadata field), then format the citation like this: https://www.youtube.com/watch?v=video_id&t=95sstart_time -> where start_time is formatted like 300s
+"""
+TA_CONTEXT_PROMPT = """Thought: I should search for information related to the users query.
+Action: Search database for information relevant to the user query.
+Information returned from the database: {context}
+Here's my final response to the user query: {query}
+FORMAT for response:
+{{Response}}
+{{citations from the metadata field returned from database}} -> if the citation is a youtube video (denoted by the metadata field), then format the citation like this: https://www.youtube.com/watch?v=video_id&t=95sstart_time -> where start_time is formatted like 300s
+"""
 
-Remember, if you can't respond using the provided information, you must respond: 'No Results Found' and nothing else.
+QUERY_AUGMENTER_SYSTEM_PROMPT = """
+You are specialized at editing a sentence that you are provided. The user will provide you a sentence 
+or a phrase and your job is to return 3 sentences or phrases that are similar to the provided sentence or phrase.
+If the sentence appears unclear or confusing, do your best to make it clear and understandable, but don't not dialogue 
+with the user, you must only respond with the 3 similar sentences or phrases and the original sentence or phrase.
+Please do not respond with anything other than the original sentence or phrase and the 3 similar sentences or phrases.
+Remember, in order to be successful, do not respond with anything other than the original sentence or phrase and the 3 similar sentences or phrases.
 """
 
 TA_SYSTEM_PROMPT = (
@@ -23,14 +49,8 @@ TA_SYSTEM_PROMPT = (
     "To do this, you need to ask questions to understand the user as best as possible. "
     "This will allow you to better understand who the person is and what they need help with. "
     "Finally, you MUST use markdown format when you respond to the user."
+    "Please ensure all links are clickable using markdown format."
 )
-TA_CONTEXT_PROMPT = """I found some information to help you. Please read it and use it to help you respond.
-
-Information: {context}
-Query: {question}
-
-Remember, if you can't respond using the provided information, you must respond: 'No Results Found' and nothing else.
-"""
 
 MODEL_CONTEXT_WINDOW = 3800
 MAX_TOKENS_FOR_RESPONSE = 800
@@ -102,12 +122,15 @@ class GPTTurboChatSession(BaseModel):
             return result
         else:
             raise StopIteration
-    
+
+    def __reversed__(self):
+        return reversed(self.messages)
+
     def add_message(self, messages: Union[Sequence[GPTTurboChat], GPTTurboChat]) -> None:
         """Add a message to the chat session and return a new chat session model"""
         if isinstance(messages, GPTTurboChat):
             messages = (messages,)
-        new_messages = self.messages + tuple(messages)
+        new_messages = self.messages + messages
         self.messages = new_messages
 
     def messages_exist(self) -> bool:
@@ -123,6 +146,13 @@ class GPTTurboChatSession(BaseModel):
         """Get the content of the last message."""
         return self.messages[-1].content
 
+    def remove_last_ai_message(self) -> None:
+        """Remove the last AI message."""
+        for i, message in enumerate(reversed(self.messages)):
+            if message.role != Role.USER:
+                self.messages = self.messages[: -(i + 1)]
+                break
+
 def count_tokens(string: str) -> int:
     """
     Get the token count of a string.
@@ -208,6 +238,7 @@ def truncate_chat_session(
 class LLMType(str, Enum):
     SEARCH_ENGINE = "search_engine"
     TA = "teaching_assistant"
+    QUERY_AUGMENTATION = "query_augmentation"
 
 class LLMWrapper:
     def __init__(self, api_key: str, llm_type: LLMType):
@@ -218,20 +249,23 @@ class LLMWrapper:
         elif llm_type == LLMType.TA:
             self._system_prompt = TA_SYSTEM_PROMPT
             self._context_prompt = TA_CONTEXT_PROMPT
+        elif llm_type == LLMType.QUERY_AUGMENTATION:
+            self._system_prompt = QUERY_AUGMENTER_SYSTEM_PROMPT
 
     def get_llm_response(self, chat_session: Union[GPTTurboChatSession, GPTTurboChat], context: Optional[Union[str, list[str]]]=None) -> Union[GPTTurboChatSession, GPTTurboChat]:
         # construct a series of messages using the conversation history
         prompt_is_instance_of_chat = isinstance(chat_session, GPTTurboChat)
         if prompt_is_instance_of_chat:
             chat_session = GPTTurboChatSession(messages=(chat_session,))
-
+        logger.info(isinstance(chat_session, GPTTurboChatSession))
+        logger.info(f"Chat session: {chat_session}")
         system_token_count = count_tokens(self._system_prompt)
         token_context_window = MODEL_CONTEXT_WINDOW
         chat_session = truncate_chat_session(chat_session, system_token_count, MAX_TOKENS_FOR_RESPONSE, token_context_window)
         if context and chat_session.messages_exist():
             context = context if isinstance(context, str) else "\n".join(context)
-            final_prompt = self._context_prompt.format(context=context, question=chat_session.last_message_content())
-            chat_session.replace_last_message(GPTTurboChat(content=final_prompt, role=Role.USER))
+            final_prompt = self._context_prompt.format(context=context, query=chat_session.last_message_content())
+            chat_session.add_message(GPTTurboChat(content=final_prompt, role=Role.ASSISTANT))
         prompt_messages = [
             {"role": Role.SYSTEM.value, "content": self._system_prompt}
         ]
@@ -240,6 +274,7 @@ class LLMWrapper:
             user_tokens_for_request += chat.token_count
             prompt_messages.append(chat.dict(exclude={"token_count"}))
 
+        logger.info(f"Prompt messages: {prompt_messages}")
         response = openai.ChatCompletion.create(
             model="gpt-3.5-turbo",
             api_key=self._api_key,
@@ -247,6 +282,8 @@ class LLMWrapper:
         )
         message = response.choices[0].message.content
         completion_tokens = response.usage.completion_tokens
+        if context and chat_session.messages_exist():
+            chat_session.remove_last_ai_message()
         chat_session.add_message(GPTTurboChat(
             role=Role.ASSISTANT,
             content=message,

taai_demo/plugin/semantic_search.py

@@ -46,10 +46,10 @@ class SemanticSearchWrapper:
         #     return get_semantic_snippets(obj.file_path)
         if isinstance(obj, TextbookPDF):
             self.object_type = "textbook"
-            return get_fix_length_snippets(obj.file_path, 1000)
+            return get_fix_length_snippets(obj.file_path, 500)
         if isinstance(obj, YoutubeVideo):
             self.object_type = "youtube"
-            return get_yt_transcript_chunks(obj.url)
+            return get_yt_transcript_chunks(obj.url, 400)
         print("I SHOULDN'T BE HERE")
 
     def insert_pdf(self, buffer) -> None:
@@ -79,13 +79,19 @@ class SemanticSearchWrapper:
             vector = Vector(vector=vector, metadata={}, text=query)
             ids = self._pinecone_index.search(
                 query=vector,
-                top_k=5,
+                top_k=25,
             )
             text = []
+            # keep the first 10 and last 5
+            logger.info(f"Found {len(ids)} relevant snippets")
+            ids = ids[:10] + ids[-3:]
+            ids = list(set(ids))
+            logger.info(f"Found {len(ids)} unique relevant snippets")
             for vector_id in ids:
                 for vector_text_mapping in self.vector_text_mappings or []: # i know this is not efficient but it's a demo and won't be slow with the amount of data we have
                     if vector_id == vector_text_mapping.vector_id:
-                        text.append(vector_text_mapping.text)
+                        text.append(f"Info: {vector_text_mapping.text}\nmetadata: {vector_text_mapping.metadata}\n")
+            logger.info(f"Relevant context: {''.join(text) if text else 'No relevant context found'}")
             return text
 
     @classmethod

taai_demo/plugin/vector_db.py

@@ -12,6 +12,7 @@ class Vector(NamedTuple):
 class VectorIdTextMapping(BaseModel):
     vector_id: str
     text: str
+    metadata: Dict[str, str]
 
 class QueryMatch(TypedDict):
     id: str
@@ -37,7 +38,7 @@ class PineconeIndex:
                 vector.vector,
                 vector.metadata,
             )
-            vector_mappings.append(VectorIdTextMapping(vector_id=str(i), text=vector.text))
+            vector_mappings.append(VectorIdTextMapping(vector_id=str(i), text=vector.text, metadata=vector.metadata))
             pinecone_vectors.append(pinecone_vector)
         try:
             self.index.upsert(vectors=pinecone_vectors, batch_size=20)

taai_demo/plugin/yt_transcript_chunker.py

@@ -19,7 +19,12 @@ def get_yt_transcript_chunks(url, max_length=500) -> Sequence[SemanticSnippet]:
         if len(chunk) > max_length or i == len(transcript) - 1:
             current_chunk.pop()
             chunk = " ".join(map(lambda x: x["text"], current_chunk))
-            chunks.append(SemanticSnippet(text=chunk, metadata={"start": current_chunk[0]["start"], "end": current_chunk[-1]["end"]}))
+            metadata = {
+                "start": current_chunk[0]["start"],
+                "end": current_chunk[-1]["end"],
+                "video_id": query["v"][0],
+            }
+            chunks.append(SemanticSnippet(text=chunk, metadata=metadata))
             current_chunk = [current_chunk[-1], snippet]
 
     return chunks

taai_demo/ui/chat_page.py

@@ -1,20 +1,26 @@
+from loguru import logger
 import streamlit as st
 from streamlit_chat import message as st_message
-from taai_demo.plugin.chat_wrapper import ChatWrapper
+from taai_demo.plugin.llm import LLMWrapper,LLMType, GPTTurboChatSession, GPTTurboChat, Role
+from taai_demo.plugin.semantic_search import SemanticSearchWrapper
+from taai_demo.settings import Settings
 
 
 def chat_page():
+    if 'message' not in st.session_state:
+        st.session_state.message = ""
+    entry: GPTTurboChat
+    for i, entry in enumerate(st.session_state.chat_history):
+        st_message(entry.content, is_user=entry.role == Role.USER.value, key=f"{i}-message", seed=st.session_state.message_seed)
     message = st.text_input("Chat with ai")
-    if message:
-        st.session_state.chat_history.append({"message": message, "role": "user"})
-
-        settings = st.session_state.settings
-        semantic_search_wrapper = st.session_state.semantic_search_wrapper
-        chat_wrapper = ChatWrapper.from_settings(settings)
+    if message != st.session_state.message:
+        chat_history: GPTTurboChatSession = st.session_state.chat_history
+        chat_history.add_message(GPTTurboChat(content=message, role="user"))
+        settings: Settings = st.session_state.settings
+        semantic_search_wrapper: SemanticSearchWrapper = st.session_state.semantic_search_wrapper
+        chat_wrapper = LLMWrapper.from_settings(settings, LLMType.TA)
 
         context = semantic_search_wrapper.search(message)
-        response = chat_wrapper.get_openai_response(st.session_state.chat_history, context)
-
-        st.session_state.chat_history.append({"message": response, "role": "assistant"})
-    for i, entry in enumerate(st.session_state.chat_history):
-        st_message(entry["message"], is_user=entry["role"]=="user", key=f"{i}-message", seed=st.session_state.message_seed)
+        st.session_state.chat_history = chat_wrapper.get_llm_response(chat_history, context)
+        st.session_state.message = message
+        st.experimental_rerun()

taai_demo/ui/search_page.py

@@ -1,11 +1,20 @@
 import streamlit as st
 from taai_demo.plugin.semantic_search import SemanticSearchWrapper
+from taai_demo.plugin.llm import GPTTurboChat, LLMWrapper, LLMType, Role
+from taai_demo.settings import Settings
 
 
 def search_page():
     query = st.text_input("Search class materials", key="search_input")
     if query:
-        semantic_search_wrapper = st.session_state.semantic_search_wrapper
-        result = semantic_search_wrapper.search(query)
-        for text in result:
-            st.markdown(text)
+        # this doesn't work super well
+        # llm_query_augmenter = LLMWrapper.from_settings(Settings(), LLMType.QUERY_AUGMENTATION)
+        # query = llm_query_augmenter.get_llm_response(GPTTurboChat(content=query, role=Role.USER))
+        # assert isinstance(query, GPTTurboChat)
+        content = query
+        semantic_search_wrapper: SemanticSearchWrapper = st.session_state.semantic_search_wrapper
+        result = semantic_search_wrapper.search(content)
+        llm = LLMWrapper.from_settings(Settings(), LLMType.SEARCH_ENGINE)
+        llm_result = llm.get_llm_response(GPTTurboChat(content=content, role=Role.USER), result)
+        assert isinstance(llm_result, GPTTurboChat)
+        st.markdown(llm_result.content)

taai_demo/ui/upload_page.py

@@ -1,11 +1,12 @@
 import streamlit as st
 from taai_demo.ui.state_management import transition_state
 from taai_demo.plugin.semantic_search import SemanticSearchWrapper
+from taai_demo.plugin.llm import GPTTurboChatSession
 from taai_demo.settings import Settings
 
 
 def upload_page():
-    st.session_state.chat_history = []
+    st.session_state.chat_history = GPTTurboChatSession()
     uploaded_file = st.file_uploader("Upload PDF")
     settings: Settings = st.session_state.settings
     semantic_search_wrapper = SemanticSearchWrapper.from_settings(settings)