Преглед изворни кода

feat: add dry run (#47)

This commit adds support to simulate the semantic search and see what the prompt will look lile.
cachho пре 2 година
родитељ
комит
f8e5ccd007
2 измењених фајлова са 40 додато и 0 уклоњено
  1. 23 0
      README.md
  2. 17 0
      embedchain/embedchain.py

+ 23 - 0
README.md

@@ -207,6 +207,29 @@ print(naval_chat_bot.query("What unique capacity does Naval argue humans possess
 
 * If you want to add any other format, please create an [issue](https://github.com/embedchain/embedchain/issues) and we will add it to the list of supported formats.
 
+## Testing
+
+Before you consume valueable tokens, you should make sure that the embedding you have done works and that it's receiving the correct document from the database. 
+
+For this you can use the `dry_run` method.
+
+Following the example above, add this to your script:
+
+```python
+print(naval_chat_bot.dry_run('Can you tell me who Naval Ravikant is?'))
+
+'''
+Use the following pieces of context to answer the query at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
+        Q: Who is Naval Ravikant?
+A: Naval Ravikant is an Indian-American entrepreneur and investor.
+        Query: Can you tell me who Naval Ravikant is?
+        Helpful Answer:
+'''
+```
+*The embedding is confirmed to work as expected. It returns the right document, even if the question is asked slightly different. No prompt tokens have been consumed.*
+
+**The dry run will still consume tokens to embed your query, but it is only ~1/15 of the prompt.**
+
 # How does it work?
 
 Creating a chat bot over any dataset needs the following steps to happen

+ 17 - 0
embedchain/embedchain.py

@@ -219,6 +219,22 @@ class EmbedChain:
         prompt = self.generate_prompt(input_query, context)
         answer = self.get_answer_from_llm(prompt)
         return answer
+    
+    def dry_run(self, input_query):
+        """
+        A dry run does everything except send the resulting prompt to
+        the LLM. The purpose is to test the prompt, not the response.
+        You can use it to test your prompt, including the context provided
+        by the vector database's doc retrieval.
+        The only thing the dry run does not consider is the cut-off due to
+        the `max_tokens` parameter.
+
+        :param input_query: The query to use.
+        :return: The prompt that would be sent to the LLM
+        """
+        context = self.retrieve_from_database(input_query)
+        prompt = self.generate_prompt(input_query, context)
+        return prompt
 
 
 class App(EmbedChain):
@@ -228,6 +244,7 @@ class App(EmbedChain):
 
     adds(data_type, url): adds the data from the given URL to the vector db.
     query(query): finds answer to the given query using vector database and LLM.
+    dry_run(query): test your prompt without consuming tokens.
     """
 
     def __int__(self, db=None, ef=None):