Skip to content

Added changes for try catch in al app retrieval #451

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jul 30, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions pebblo/app/models/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -277,3 +277,9 @@ class RetrievalResponse(BaseModel):
class PromptResponseModel(BaseModel):
retrieval_data: Union[RetrievalResponse, None] = None
message: Optional[str] = None


class PromptGovResponseModel(BaseModel):
entities: dict
entityCount: int
message: Optional[str] = None
108 changes: 64 additions & 44 deletions pebblo/app/service/local_ui_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -162,35 +162,40 @@ def get_prompt_details(
entity_detected = prompt.get("entities")

for key, value in entity_detected.items():
if key in prompt_details[app_name]:
# If the entity type already exists in prompt_details, update the existing entry
prompt_details[app_name][key]["total_prompts"] += 1
prompt_details[app_name][key]["total_entity_count"] += value

# Get the user name from the retrieval data
user_name = retrieval.get("user", "")

if (
user_name
and user_name not in prompt_details[app_name][key]["users"]
):
# Add new user to the user list if not already present
prompt_details[app_name][key]["users"].append(user_name)
prompt_details[app_name][key]["total_users"] += 1

else:
# If the entity type does not exist in prompt_details, create a new entry
prompt_details[app_name][key] = {
"total_prompts": 1,
"total_entity_count": value,
"users": [],
"total_users": 0,
}

user_name = retrieval.get("user", "")
if user_name:
prompt_details[app_name][key]["users"].append(user_name)
prompt_details[app_name][key]["total_users"] = 1
try:
if key in prompt_details[app_name]:
# If the entity type already exists in prompt_details, update the existing entry
prompt_details[app_name][key]["total_prompts"] += 1
prompt_details[app_name][key]["total_entity_count"] += value

# Get the user name from the retrieval data
user_name = retrieval.get("user", "")

if (
user_name
and user_name not in prompt_details[app_name][key]["users"]
):
# Add new user to the user list if not already present
prompt_details[app_name][key]["users"].append(user_name)
prompt_details[app_name][key]["total_users"] += 1

else:
# If the entity type does not exist in prompt_details, create a new entry
prompt_details[app_name][key] = {
"total_prompts": 1,
"total_entity_count": value,
"users": [],
"total_users": 0,
}

user_name = retrieval.get("user", "")
if user_name:
prompt_details[app_name][key]["users"].append(user_name)
prompt_details[app_name][key]["total_users"] = 1
except Exception as ex:
logger.warning(
f"[Dashboard]: Error while iterating prompts for app {app_name} Error: {ex}"
)

prompt_details[app_name] = dict(
sorted(
Expand Down Expand Up @@ -228,12 +233,17 @@ def prepare_retrieval_response(
prompt_details[app_name] = {}
# fetch total retrievals
for retrieval in app_metadata_content.get("retrievals", []):
retrieval_data = {"name": app_json.get("name")}
retrieval_data.update(retrieval)
self.total_retrievals.append(retrieval_data)
prompt_details, total_prompt_with_findings = self.get_prompt_details(
prompt_details, retrieval, app_name, total_prompt_with_findings
)
try:
retrieval_data = {"name": app_json.get("name")}
retrieval_data.update(retrieval)
self.total_retrievals.append(retrieval_data)
prompt_details, total_prompt_with_findings = self.get_prompt_details(
prompt_details, retrieval, app_name, total_prompt_with_findings
)
except Exception as ex:
logger.warning(
f"[Dashboard]: Error while iterating retrieval for app {app_name} Error: {ex}"
)

# fetch active users per app
active_users = self.get_active_users(app_metadata_content.get("retrievals", []))
Expand Down Expand Up @@ -320,15 +330,25 @@ def get_all_apps_details(self):
)
final_prompt_details = []
for key, value in prompt_details.items():
for k1, v1 in value.items():
prompt_dict = {}
prompt_dict = {"app_name": key}
prompt_dict["entity_name"] = k1
prompt_dict["total_prompts"] = v1["total_prompts"]
prompt_dict["total_entity_count"] = v1["total_entity_count"]
prompt_dict["users"] = v1["users"]
prompt_dict["total_users"] = v1["total_users"]
final_prompt_details.append(prompt_dict)
try:
for k1, v1 in value.items():
try:
prompt_dict = {}
prompt_dict = {"app_name": key}
prompt_dict["entity_name"] = k1
prompt_dict["total_prompts"] = v1["total_prompts"]
prompt_dict["total_entity_count"] = v1["total_entity_count"]
prompt_dict["users"] = v1["users"]
prompt_dict["total_users"] = v1["total_users"]
final_prompt_details.append(prompt_dict)
except Exception as ex:
logger.warning(
f"[Dashboard]: Error in iterating key value pair in retrieval app list. Error: {ex}"
)
except Exception as ex:
logger.warning(
f"[Dashboard]: Error in iterating prompt details for all retrieval apps: {ex}"
)

# Sort loader apps
sorted_loader_apps = self._sort_loader_apps(all_loader_apps)
Expand Down
32 changes: 23 additions & 9 deletions pebblo/app/service/prompt_gov.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,10 @@

import traceback

from fastapi import HTTPException
from pydantic import ValidationError

from pebblo.app.models.models import AiDataModel
from pebblo.app.libs.responses import PebbloJsonResponse
from pebblo.app.models.models import AiDataModel, PromptGovResponseModel
from pebblo.entity_classifier.entity_classifier import EntityClassifier
from pebblo.log import get_logger

Expand Down Expand Up @@ -64,18 +64,32 @@ def process_request(self):
doc_info = self._get_classifier_response()
logger.debug(f"Entities {doc_info.entities}")
logger.debug(f"Entity Count {doc_info.entityCount}")
return {
"message": "Prompt Gov Processed Successfully",
"entities": doc_info.entities,
"entityCount": doc_info.entityCount,
}
response = PromptGovResponseModel(
entities=doc_info.entities,
entityCount=doc_info.entityCount,
message="Prompt Governance Processed Successfully",
)
return PebbloJsonResponse.build(
body=response.dict(exclude_none=True), status_code=200
)

except ValidationError as ex:
response = PromptGovResponseModel(
entities={}, entityCount=0, message=f"Error : {str(ex)}"
)
logger.error(
f"Error in Prompt API process_request. Error:{traceback.format_exc()}"
)
raise HTTPException(status_code=400, detail=str(ex))
return PebbloJsonResponse.build(
body=response.dict(exclude_none=True), status_code=400
)
except Exception as ex:
response = PromptGovResponseModel(
entities={}, entityCount=0, message=f"Error : {str(ex)}"
)
logger.error(
f"Error in Prompt API process_request. Error:{traceback.format_exc()}"
)
raise HTTPException(status_code=500, detail=str(ex))
return PebbloJsonResponse.build(
body=response.dict(exclude_none=True), status_code=500
)
18 changes: 11 additions & 7 deletions tests/app/service/test_prompt_gov.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@

import pytest

from pebblo.app.libs.responses import PebbloJsonResponse
from pebblo.app.models.models import PromptGovResponseModel
from pebblo.app.service.prompt_gov import PromptGov


Expand All @@ -23,14 +25,16 @@ def test_process_request_success(mock_entity_classifier):
data = {"prompt": "Sachin's SSN is 222-85-4836"}
prompt_gov = PromptGov(data)
response = prompt_gov.process_request()
expected_response = PromptGovResponseModel(
entities={"us-ssn": 1},
entityCount=1,
message="Prompt Governance Processed Successfully",
)
expected_response = PebbloJsonResponse.build(
body=expected_response.dict(exclude_none=True), status_code=200
)

expected_response = {
"message": "Prompt Gov Processed Successfully",
"entities": {"us-ssn": 1},
"entityCount": 1,
}

assert response == expected_response
assert response.__dict__ == expected_response.__dict__


if __name__ == "__main__":
Expand Down