Merge pull request #39 from datamindedbe/feature/chatbot-response-script

retrieval as script

backend/generation_script.py

@@ -0,0 +1,37 @@
+# a test script to implement the RAG procedure
+# should be replaced by a lambda implementation
+
+
+# do the retrieval via bedrock
+from src.bedrock import retrieve_bedrock_items, re_reference, BedrockRetrievedItem
+from src.llm import decisions_query
+from config import BEDROCK_KNOWLEDGE_BASE_ID
+import json
+from typing import List
+
+def retrieval(query: str):
+    if len(query) < 10:
+        return {"statusCode": 400, "body": json.dumps("Vraag niet lang genoeg.")}
+
+    if len(query) > 1000:
+        return {"statusCode": 400, "body": json.dumps("Vraag te lang.")}
+
+    response = retrieve_bedrock_items(BEDROCK_KNOWLEDGE_BASE_ID, query, 10)
+    return response
+
+# do the generation
+def generate(query: str, decisions: List[BedrockRetrievedItem]):
+    response = decisions_query(query, decisions)
+    return response
+
+
+
+# run
+if __name__ == "__main__":
+   query = "wat heeft de regering gedaan om de economie te stimuleren?"
+   context = retrieval(query)
+   response = generate(query, context)
+   response, used_decisions = re_reference(response, context)
+   print(response)
+   print('---')
+   print([item.s3_uri for item in used_decisions])

backend/src/bedrock.py

@@ -3,7 +3,8 @@
 
 import boto3
 import json
-from typing import Optional
+import re
+from typing import Optional, List, Tuple
 from dataclasses import dataclass
 from functools import lru_cache
 from config import AWS_PROFILE_NAME, BEDROCK_REGION, BEDROCK_BUCKET
@@ -82,3 +83,23 @@ def retrieve_bedrock_items(knowledge_base_id: str, query: str, n_results: int =
             break 
 
     return results
+
+# may not be the most sensible place for this function to live but it is a first pass
+def re_reference(response: str, decisions: List[BedrockRetrievedItem]) -> Tuple[str, List[BedrockRetrievedItem]]:
+    used_decision_references = []
+    for item in re.finditer("\[[0-9]*\]", response):
+        ref = int(item.group(0).strip("[]"))
+        if ref not in used_decision_references:
+            used_decision_references.append(ref)
+    print(used_decision_references)
+
+    # reduce the decisions to only those that are used
+    used_decisions_correct_order  = []
+    for ref in used_decision_references:
+        used_decisions_correct_order.append(decisions[ref])
+
+    # reorder the references by the order of occurence
+    for new_ref, old_ref in enumerate(used_decision_references):
+        response = response.replace(f"[{old_ref}]", f"[{new_ref}]")
+
+    return response, used_decisions_correct_order

backend/src/llm.py

@@ -3,6 +3,8 @@
 import tiktoken
 from config import OPENAI_API_KEY, OPENAI_MODEL_NAME
 from src.vector_store import VectorDBItem
+from src.bedrock import BedrockRetrievedItem
+# TODO: Bedrock items and vector db items need to combined into a single type
 
 client = openai.OpenAI(
     api_key=OPENAI_API_KEY
@@ -75,16 +77,22 @@ def extended_message(message:str)->Optional[str]:
 
     return get_response(prompt, model_overide="gpt-3.5-turbo")
 
-def decisions_query(query:str, matching_decisions:list[VectorDBItem])->str:
+
+
+def decisions_query(query:str, decisions:list[BedrockRetrievedItem])->str:
     # this function will take a given query and matching decisions and generate a response
-    flat_decisions = [f"{index} [url: {item.metadata['source_url']}] - {item.text}" for index, item in enumerate(matching_decisions)]
+    flat_decisions = [f"{index} : {item.text}" for index, item in enumerate(decisions)]
     flat_decisions = "\n".join(flat_decisions)
     prompt = f"""
-                Beantwoord de gegeven vraag met behulp van alleen de informatie in het onderstaande inhoudsgedeelte
-                inhoud: {flat_decisions}
+                Beantwoord de gegeven vraag uitsluitend met behulp van de informatie in het onderstaande gedeelte over regeringsbeslissingen.
+                reageer op een gemoedelijke manier, niet met een lijst
+                geef het indexnummer van de gebruikte referentie op tussen enkelvoudige vierkante haken, zoals [1].
+
+                regeringsbeslissingen: {flat_decisions}
                 vraag:{query}
                 Jou antwoord:
     """
+    print(f"full prompt contains: {token_count(prompt)} tokens")
     return get_response(prompt)