Merge pull request #43 from basedosdados/refactor/viz-agent

refactor: `VizAgent` and `RouterAgent`

Author: vrtornisiello

chatbot/agents/prompts.py

@@ -4,25 +4,41 @@
 from pydantic import BaseModel, Field
 
 
-class BaseItem(BaseModel):
+class Item(BaseModel):
+    id: uuid.UUID = Field(default_factory=uuid.uuid4)
     content: Any = Field(default=None)
+
+class ItemRemove(BaseModel):
     id: uuid.UUID
 
-class Item(BaseItem):
-    id: uuid.UUID = Field(default_factory=uuid.uuid4)
+class ChatTurn(BaseModel):
+    id: int # easier for the LLMs to parse
+    user_question: str
+    ai_response: str
+    data: list[Item]
 
-class ItemRemove(BaseItem):
-    ...
+class ChatTurnRemove(BaseModel):
+    id: int # easier for the LLMs to parse
 
-def add_item(existing: Item|list[Item], update: BaseItem|list[BaseItem]) -> list[Item]:
-    """Merges two lists of Items, deleting or updating Items by ID.
+ItemUpdate = Item | ItemRemove | list[Item|ItemRemove]
+
+def add_item(existing: Item | list[Item], update: ItemUpdate) -> list[Item]:
+    """Merge Items into an existing list, replacing or removing entries by ID.
 
     Args:
-        existing (list[Item]): A base list of Item objects
-        update (list[BaseItem]): A list of Item or ItemRemove objects
+        existing (Item | list[Item]):
+            The current `Item` or list of `Item` objects to be merged into.
+        update (ItemUpdate):
+            An `Item`, `ItemRemove`, or a list containing a mix of `Item` and
+            `ItemRemove` objects that specify changes to apply.
 
     Returns:
-        list[Item]: The merged list
+        list[Item]:
+            The merged list of `Item` objects after applying the updates and deletions.
+
+    Raises:
+        ValueError:
+            If an `ItemRemove` references an ID that does not exist in `existing`.
     """
     if not isinstance(existing, list):
         existing = [existing]
@@ -31,7 +47,7 @@ def add_item(existing: Item|list[Item], update: BaseItem|list[BaseItem]) -> list
         update = [update]
 
     # copying to avoid modifying the existing list
-    merged: list[Item] = existing.copy() or []
+    merged: list[Item] = existing.copy()
 
     ids_to_idx = {item.id: idx for idx, item in enumerate(existing)}
 
@@ -55,3 +71,34 @@ def add_item(existing: Item|list[Item], update: BaseItem|list[BaseItem]) -> list
     merged = [item for item in merged if item.id not in ids_to_remove]
 
     return merged
+
+def add_chat_turn(existing: dict[int, ChatTurn], update: dict[int, ChatTurn|ChatTurnRemove]) -> dict[int, ChatTurn]:
+    """Merge chat turns into an existing dictionary, replacing or removing entries by key.
+
+    Args:
+        existing (dict[int, ChatTurn]):
+            The current dictionary of chat turns to be updated.
+        update (dict[int, ChatTurn | ChatTurnRemove]):
+            The dictionary containing additions, replacements, or deletions to apply.
+
+    Returns:
+        dict[int, ChatTurn]:
+            The updated dictionary of chat turns.
+
+    Raises:
+        ValueError:
+            If a `ChatTurnRemove` references a key not present in `existing`.
+    """
+    merged = existing.copy()
+
+    for k, v in update.items():
+        if isinstance(v, ChatTurnRemove):
+            if k not in merged:
+                raise ValueError(
+                    f"Attempted to delete non-existent chat turn '{k}'"
+                )
+            del merged[k]
+        else:
+            merged[k] = v
+
+    return merged