Working self route

i-dot-ai · Sep 6, 2024 · 48a0f56 · 48a0f56
1 parent 38f6987
commit 48a0f56
Show file tree

Hide file tree

Showing 7 changed files with 378 additions and 383 deletions.
diff --git a/.vscode/redbox.code-workspace b/.vscode/redbox.code-workspace
@@ -4,10 +4,6 @@
 			"name": "Django App",
 			"path": "../django_app"
 		},
-		{
-			"name": "Core API",
-			"path": "../core-api"
-		},
 		{
 			"name": "Redbox Core",
 			"path": "../redbox-core"

diff --git a/django_app/poetry.lock b/django_app/poetry.lock
diff --git a/django_app/pyproject.toml b/django_app/pyproject.toml
@@ -77,4 +77,8 @@ build-backend = "poetry.core.masonry.api"
 DJANGO_SETTINGS_MODULE = "redbox_app.settings"
 testpaths = "tests"
 norecursedirs = "tests/test_ai.py"
-
+env_override_existing_values = 1
+env_files = [
+    ".env.test",
+    ".env"
+]
diff --git a/django_app/redbox_app/redbox_core/consumers.py b/django_app/redbox_app/redbox_core/consumers.py
@@ -16,8 +16,7 @@
 
 from redbox import Redbox
 from redbox.models import Settings
-from redbox.models.chain import ChainChatMessage, RedboxQuery, RedboxState
-from redbox.models.chat import MetadataDetail
+from redbox.models.chain import ChainChatMessage, RedboxQuery, RedboxState, RequestMetadata, metadata_reducer
 from redbox_app.redbox_core import error_messages
 from redbox_app.redbox_core.models import (
     Chat,
@@ -50,15 +49,14 @@ class ChatConsumer(AsyncWebsocketConsumer):
     full_reply: ClassVar = []
     citations: ClassVar = []
     route = None
-    metadata: MetadataDetail = MetadataDetail()
+    metadata: RequestMetadata = RequestMetadata()
     redbox = Redbox(env=Settings(), debug=True)
 
     async def receive(self, text_data=None, bytes_data=None):
         """Receive & respond to message from browser websocket."""
         self.full_reply = []
         self.citations = []
         self.route = None
-        self.metadata = MetadataDetail()
 
         data = json.loads(text_data or bytes_data)
         logger.debug("received %s from browser", data)
@@ -142,7 +140,7 @@ def save_message(
         role: ChatRoleEnum,
         sources: Sequence[tuple[File, Document]] | None = None,
         selected_files: Sequence[File] | None = None,
-        metadata: MetadataDetail | None = None,
+        metadata: RequestMetadata | None = None,
         route: str | None = None,
     ) -> ChatMessage:
         chat_message = ChatMessage(chat=session, text=user_message_text, role=role, route=route)
@@ -194,11 +192,7 @@ async def handle_route(self, response: str) -> str:
         self.route = response
 
     async def handle_metadata(self, response: dict):
-        metadata_detail = MetadataDetail.parse_obj(response)
-        for model, token_count in metadata_detail.input_tokens.items():
-            self.metadata.input_tokens[model] = self.metadata.input_tokens.get(model, 0) + token_count
-        for model, token_count in metadata_detail.output_tokens.items():
-            self.metadata.output_tokens[model] = self.metadata.output_tokens.get(model, 0) + token_count
+        self.metadata = metadata_reducer(self.metadata, RequestMetadata.model_validate(response))
 
     async def handle_documents(self, response: list[Document]):
         s3_keys = [doc.metadata["file_name"] for doc in response]

diff --git a/django_app/tests/test_consumers.py b/django_app/tests/test_consumers.py
@@ -17,7 +17,7 @@
 from websockets.legacy.client import Connect
 
 from redbox.graph.root import FINAL_RESPONSE_TAG, ROUTE_NAME_TAG, SOURCE_DOCUMENTS_TAG
-from redbox.models.chat import MetadataDetail
+from redbox.models.chain import RequestMetadata, LLMCallMetadata
 from redbox_app.redbox_core import error_messages
 from redbox_app.redbox_core.consumers import ChatConsumer
 from redbox_app.redbox_core.models import Chat, ChatMessage, ChatMessageTokenUse, ChatRoleEnum, File, User
@@ -508,7 +508,11 @@ def mocked_connect(uploaded_file: File) -> Connect:
         {
             "event": "on_custom_event",
             "name": "on_metadata_generation",
-            "data": MetadataDetail(input_tokens={"gpt-4o": 123}, output_tokens={"gpt-4o": 1000}),
+            "data": RequestMetadata(
+                llm_calls=[LLMCallMetadata(model_name="gpt-4o", input_tokens=123, output_tokens=1000)],
+                selected_files_total_tokens=1000,
+                number_of_selected_files=1
+            ),
         },
     ]
 

diff --git a/redbox-core/redbox/models/chat.py b/redbox-core/redbox/models/chat.py
@@ -11,9 +11,4 @@ class ChatRoute(StrEnum):
 
 
 class ErrorRoute(StrEnum):
-    files_too_large = "error/files_too_large"
-
-
-class MetadataDetail(BaseModel):
-    input_tokens: dict[str, int] = Field(default_factory=dict)
-    output_tokens: dict[str, int] = Field(default_factory=dict)
+    files_too_large = "error/files_too_large"
diff --git a/redbox-core/redbox/retriever/retrievers.py b/redbox-core/redbox/retriever/retrievers.py
@@ -27,7 +27,7 @@ def hit_to_doc(hit: dict[str, Any]) -> Document:
         "uuid": hit["_id"],
     }
     return Document(
-        page_content=source.get("text"), metadata={k: v for k, v in c_meta.items() if v is not None} | source["metadata"]
+        page_content=source.get("text", ""), metadata={k: v for k, v in c_meta.items() if v is not None} | source["metadata"]
     )