Skip to content

Commit

Permalink
AAP-30619: migrate PlaybookGeneration to the new Schema1 object
Browse files Browse the repository at this point in the history 
Use an object defined in ansible_ai_connect/ai/api/telemetry/schema1.py to prepare the Schema1 payload.
  • Loading branch information
@goneri
goneri committed Jan 8, 2025
1 parent 48acf18 commit 85d64f1
Show file tree
Hide file tree
Showing 3 changed files with 69 additions and 173 deletions.
11 changes: 11 additions & 0 deletions ansible_ai_connect/ai/api/telemetry/schema1.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,17 @@ class ExplainPlaybookEvent(Schema1Event):
explanationId: str = field(validator=validators.instance_of(str), converter=str, default="")


@define
class GenerationPlaybookEvent(Schema1Event):
event_name: str = "codegenPlaybook"
playbook_length: int = field(validator=validators.instance_of(int), default=0)
generationId: str = field(validator=validators.instance_of(str), converter=str, default="")
wizardId: str = field(validator=validators.instance_of(str), converter=str, default="")
create_outline: bool = field(
validator=validators.instance_of(bool), converter=bool, default=False
)


@define
class ChatBotResponseDocsReferences:
docs_url: str = field(validator=validators.instance_of(str), converter=str, default="")
Expand Down
16 changes: 15 additions & 1 deletion ansible_ai_connect/ai/api/telemetry/test_schema1.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,12 @@
from ansible_ai_connect.test_utils import WisdomServiceAPITestCaseBaseOIDC
from ansible_ai_connect.users.models import Plan

from .schema1 import ExplainPlaybookEvent, OneClickTrialStartedEvent, Schema1Event
from .schema1 import (
ExplainPlaybookEvent,
GenerationPlaybookEvent,
OneClickTrialStartedEvent,
Schema1Event,
)


@override_settings(AUTHZ_BACKEND_TYPE="dummy")
Expand Down Expand Up @@ -107,3 +112,12 @@ class TestExplainPlaybookEvent(WisdomServiceAPITestCaseBaseOIDC):
def test_base(self):
event1 = ExplainPlaybookEvent()
self.assertEqual(event1.event_name, "explainPlaybook")


@override_settings(AUTHZ_BACKEND_TYPE="dummy")
@override_settings(WCA_SECRET_DUMMY_SECRETS="1981:valid")
class TestGenerationPlaybookEvent(WisdomServiceAPITestCaseBaseOIDC):
def test_base(self):
event1 = GenerationPlaybookEvent()
self.assertEqual(event1.event_name, "codegenPlaybook")
self.assertFalse(event1.create_outline)
215 changes: 43 additions & 172 deletions ansible_ai_connect/ai/api/views.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -797,13 +797,14 @@ def post(self, request) -> Response:
)


class GenerationPlaybook(APIView):
class GenerationPlaybook(AACSAPIView):
"""
Returns a playbook based on a text input.
"""

permission_classes = PERMISSIONS_MAP.get(settings.DEPLOYMENT_MODE)
required_scopes = ["read", "write"]
schema1_event = schema1.GenerationPlaybookEvent

throttle_cache_key_suffix = "_generation_playbook"

Expand All @@ -820,195 +821,65 @@ class GenerationPlaybook(APIView):
summary="Inline code suggestions",
)
def post(self, request) -> Response:
exception = None
generation_id = None
wizard_id = None
duration = None
create_outline = None
anonymized_playbook = ""
playbook = ""
request_serializer = GenerationPlaybookRequestSerializer(data=request.data)
answer = {}
model_id = ""
try:
request_serializer.is_valid(raise_exception=True)
generation_id = str(request_serializer.validated_data.get("generationId", ""))
create_outline = request_serializer.validated_data["createOutline"]
outline = str(request_serializer.validated_data.get("outline", ""))
text = request_serializer.validated_data["text"]
custom_prompt = str(request_serializer.validated_data.get("customPrompt", ""))
wizard_id = str(request_serializer.validated_data.get("wizardId", ""))
model_id = str(request_serializer.validated_data.get("model", ""))

llm: ModelPipelinePlaybookGeneration = apps.get_app_config("ai").get_model_pipeline(
ModelPipelinePlaybookGeneration
)
start_time = time.time()
playbook, outline, warnings = llm.invoke(
PlaybookGenerationParameters.init(
request=request,
text=text,
custom_prompt=custom_prompt,
create_outline=create_outline,
outline=outline,
generation_id=generation_id,
model_id=model_id,
)
)
duration = round((time.time() - start_time) * 1000, 2)

# Anonymize responses
# Anonymized in the View to be consistent with where Completions are anonymized
anonymized_playbook = anonymizer.anonymize_struct(
playbook, value_template=Template("{{ _${variable_name}_ }}")
)
anonymized_outline = anonymizer.anonymize_struct(
outline, value_template=Template("{{ _${variable_name}_ }}")
)

answer = {
"playbook": anonymized_playbook,
"outline": anonymized_outline,
"warnings": warnings,
"format": "plaintext",
"generationId": generation_id,
}

except WcaBadRequest as e:
exception = e
logger.exception(f"bad request for playbook generation {generation_id}")
raise WcaBadRequestException(cause=e)

except WcaInvalidModelId as e:
exception = e
logger.exception(f"WCA Model ID is invalid for playbook generation {generation_id}")
raise WcaInvalidModelIdException(cause=e)

except WcaKeyNotFound as e:
exception = e
logger.exception(
f"A WCA Api Key was expected but not found for "
f"playbook generation {generation_id}"
)
raise WcaKeyNotFoundException(cause=e)

except WcaModelIdNotFound as e:
exception = e
logger.exception(
f"A WCA Model ID was expected but not found for "
f"playbook generation {generation_id}"
)
raise WcaModelIdNotFoundException(cause=e)

except WcaNoDefaultModelId as e:
exception = e
logger.exception(
"A default WCA Model ID was expected but not found for "
f"playbook generation {generation_id}"
)
raise WcaNoDefaultModelIdException(cause=e)

except WcaRequestIdCorrelationFailure as e:
exception = e
logger.exception(
f"WCA Request/Response GenerationId correlation failed "
f"for suggestion {generation_id}"
)
raise WcaRequestIdCorrelationFailureException(cause=e)

except WcaEmptyResponse as e:
exception = e
logger.exception(
f"WCA returned an empty response for playbook generation {generation_id}"
)
raise WcaEmptyResponseException(cause=e)

except WcaCloudflareRejection as e:
exception = e
logger.exception(
f"Cloudflare rejected the request for playbook generation {generation_id}"
)
raise WcaCloudflareRejectionException(cause=e)

except WcaHAPFilterRejection as e:
exception = e
logger.exception(
f"WCA Hate, Abuse, and Profanity filter rejected "
f"the request for playbook generation {generation_id}"
)
raise WcaHAPFilterRejectionException(cause=e)

except WcaUserTrialExpired as e:
exception = e
logger.exception(
f"User trial expired, when requesting playbook generation {generation_id}"
)
raise WcaUserTrialExpiredException(cause=e)
request_serializer = GenerationPlaybookRequestSerializer(data=request.data)
request_serializer.is_valid(raise_exception=True)
generation_id = str(request_serializer.validated_data.get("generationId", ""))
create_outline = request_serializer.validated_data["createOutline"]
outline = str(request_serializer.validated_data.get("outline", ""))
text = request_serializer.validated_data["text"]
custom_prompt = str(request_serializer.validated_data.get("customPrompt", ""))
wizard_id = str(request_serializer.validated_data.get("wizardId", ""))
model_id = str(request_serializer.validated_data.get("model", ""))

except WcaInstanceDeleted as e:
exception = e
logger.exception(
"WCA Instance has been deleted when requesting playbook generation "
f"{generation_id} for model {e.model_id}"
self.event.create_outline = create_outline
llm: ModelPipelinePlaybookGeneration = apps.get_app_config("ai").get_model_pipeline(
ModelPipelinePlaybookGeneration
)
playbook, outline, warnings = llm.invoke(
PlaybookGenerationParameters.init(
request=request,
text=text,
custom_prompt=custom_prompt,
create_outline=create_outline,
outline=outline,
generation_id=generation_id,
model_id=model_id,
)
raise WcaInstanceDeletedException(cause=e)
)
self.event.playbook_length = len(playbook)
self.event.generationId = generation_id
self.event.wizard_id = wizard_id

except Exception as exc:
exception = exc
logger.exception(f"An exception {exc.__class__} occurred during a playbook generation")
raise
# Anonymize responses
# Anonymized in the View to be consistent with where Completions are anonymized
anonymized_playbook = anonymizer.anonymize_struct(
playbook, value_template=Template("{{ _${variable_name}_ }}")
)
anonymized_outline = anonymizer.anonymize_struct(
outline, value_template=Template("{{ _${variable_name}_ }}")
)

finally:
self.write_to_segment(
request.user,
generation_id,
wizard_id,
exception,
duration,
create_outline,
playbook_length=len(anonymized_playbook),
model_id=model_id,
)
answer = {
"playbook": anonymized_playbook,
"outline": anonymized_outline,
"warnings": warnings,
"format": "plaintext",
"generationId": generation_id,
}

return Response(
answer,
status=rest_framework_status.HTTP_200_OK,
)

def write_to_segment(
self,
user,
generation_id,
wizard_id,
exception,
duration,
create_outline,
playbook_length,
model_id,
):
model_name = ""
try:
model_meta_data: MetaData = apps.get_app_config("ai").get_model_pipeline(MetaData)
model_name = model_meta_data.get_model_id(user, user.org_id, model_id)
except (WcaNoDefaultModelId, WcaModelIdNotFound, WcaSecretManagerError):
pass
event = {
"create_outline": create_outline,
"duration": duration,
"exception": exception is not None,
"generationId": generation_id,
"modelName": model_name,
"playbook_length": playbook_length,
"wizardId": wizard_id,
}
if exception:
event["response"] = (
{
"exception": str(exception),
},
)
send_segment_event(event, "codegenPlaybook", user)


class GenerationRole(APIView):
"""
Expand Down

0 comments on commit 85d64f1

Please sign in to comment.