fix: Format resource list to show relevant resource fields clearly (#25)

Author: yoomlam

app/src/app_config.py

@@ -31,7 +31,7 @@ def db_session(self) -> db.Session:
     PROMPT_VERSIONS: dict = {
         "sample_rag": "UHJvbXB0VmVyc2lvbjox",
         "extract_supports": "UHJvbXB0VmVyc2lvbjoz",
-        "generate_referrals": "UHJvbXB0VmVyc2lvbjo0",
+        "generate_referrals": "UHJvbXB0VmVyc2lvbjo2",
     }
 
 

app/src/pipelines/generate_referrals/pipeline_wrapper.py

@@ -1,6 +1,7 @@
 import json
 import logging
 from pprint import pformat
+from uuid import UUID
 
 import hayhooks
 from hayhooks import BasePipelineWrapper
@@ -9,7 +10,6 @@
 from haystack.dataclasses.chat_message import ChatMessage
 from haystack_integrations.components.generators.amazon_bedrock import AmazonBedrockChatGenerator
 from pydantic import BaseModel
-from sqlalchemy.inspection import inspect
 
 from src.app_config import config
 from src.common import haystack_utils
@@ -19,9 +19,11 @@
 
 
 class Resource(BaseModel):
-    resource_name: str
-    resource_addresses: list[str]
-    resource_phones: list[str]
+    name: str
+    addresses: list[str]
+    phones: list[str]
+    emails: list[str]
+    website: str
     description: str
     justification: str
 
@@ -50,12 +52,12 @@ def setup(self) -> None:
 
     # Called for the `generate-referrals/run` endpoint
     def run_api(self, query: str) -> dict:
-        supports_from_db = retrieve_supports_from_db()
+        supports_from_db = format_support_strings()
         response = self.pipeline.run(
             {
                 "prompt_builder": {
                     "query": query,
-                    "supports": supports_from_db,
+                    "supports": supports_from_db.values(),
                     "resource_json": resource_as_json,
                 },
             }
@@ -80,15 +82,16 @@ def run_chat_completion(self, model: str, messages: list, body: dict) -> None:
         )
 
 
-def retrieve_supports_from_db() -> list[str]:
-    all_supports: list[str] = []
+def format_support_strings() -> dict[UUID, str]:
     with config.db_session() as db_session, db_session.begin():
-        all_db_supports = db_session.query(Support).all()
-
-        for support in all_db_supports:
-            support_dict = {
-                c.key: getattr(support, c.key) for c in inspect(Support).mapper.column_attrs
-            }
-            support_as_str = json.dumps(support_dict, default=str)
-            all_supports.append(support_as_str)
-    return all_supports
+        return {
+            support.id: (
+                f"Name: {support.name}\n"
+                f"- Description: {support.description}\n"
+                f"- Addresses: {', '.join(support.addresses)}\n"
+                f"- Phones: {', '.join(support.phone_numbers)}\n"
+                f"- Website: {support.website}\n"
+                f"- Email Addresses: {', '.join(support.email_addresses)}\n"
+            )
+            for support in db_session.query(Support).all()
+        }

app/tests/src/pipelines/generate_referrals/test_generate_referrals_pipeline.py

@@ -1,35 +1,35 @@
-import json
-
 import pytest
-from sqlalchemy.inspection import inspect
 
 from src.adapters import db
 from src.db.models.support_listing import Support
-from src.pipelines.generate_referrals.pipeline_wrapper import retrieve_supports_from_db
+from src.pipelines.generate_referrals.pipeline_wrapper import format_support_strings
 from tests.src.db.models.factories import SupportFactory, SupportListingFactory
 
 
 @pytest.fixture
-def seed_supports(db_session: db.Session):
+def seed_supports(enable_factory_create, db_session: db.Session):
     # remove all pre-existing Support records
     db_session.query(Support).delete()
 
     support_listing = SupportListingFactory.create()
-    supports = []
-    for i in range(0, 3):
-        support = SupportFactory.create(support_listing=support_listing, name=f"support{i}")
-        support_as_json_str = json.dumps(
-            {c.key: getattr(support, c.key) for c in inspect(Support).mapper.column_attrs},
-            default=str,
-        )
-        supports.append(support_as_json_str)
-    return supports
-
-
-def test_retrieve_supports_from_db(enable_factory_create, seed_supports, db_session: db.Session):
-    supports_from_db = retrieve_supports_from_db()
-
-    assert len(supports_from_db) == 3
-    assert seed_supports[0] == supports_from_db[0]
-    assert seed_supports[1] == supports_from_db[1]
-    assert seed_supports[2] == supports_from_db[2]
+    return [SupportFactory.create(support_listing=support_listing) for _ in range(3)]
+
+
+def test_format_support_strings(seed_supports):
+    support_strings = format_support_strings()
+
+    assert len(support_strings) == len(seed_supports)
+
+    db_supports = {support.id: support for support in seed_supports}
+    for k, v in support_strings.items():
+        support = db_supports[k]
+        # Verify that support data is somewhere in the resulting strings
+        assert support.name in v
+        assert support.description in v
+        for address in support.addresses:
+            assert address in v
+        for phone_number in support.phone_numbers:
+            assert phone_number in v
+        assert support.website in v
+        for email_address in support.email_addresses:
+            assert email_address in v