Add AI suggestion endpoints (#1519)

skyvern/forge/sdk/db/client.py

@@ -14,6 +14,7 @@ from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType, TaskType
 from skyvern.forge.sdk.db.exceptions import NotFoundError
 from skyvern.forge.sdk.db.models import (
     ActionModel,
+    AISuggestionModel,
     ArtifactModel,
     AWSSecretParameterModel,
     BitwardenCreditCardDataParameterModel,
@@ -56,6 +57,7 @@ from skyvern.forge.sdk.db.utils import (
 )
 from skyvern.forge.sdk.log_artifacts import save_workflow_run_logs
 from skyvern.forge.sdk.models import Step, StepStatus
+from skyvern.forge.sdk.schemas.ai_suggestions import AISuggestion
 from skyvern.forge.sdk.schemas.observers import (
     ObserverCruise,
     ObserverCruiseStatus,
@@ -206,6 +208,7 @@ class AgentDB:
         workflow_run_block_id: str | None = None,
         observer_cruise_id: str | None = None,
         observer_thought_id: str | None = None,
+        ai_suggestion_id: str | None = None,
         organization_id: str | None = None,
     ) -> Artifact:
         try:
@@ -220,6 +223,7 @@ class AgentDB:
                     workflow_run_block_id=workflow_run_block_id,
                     observer_cruise_id=observer_cruise_id,
                     observer_thought_id=observer_thought_id,
+                    ai_suggestion_id=ai_suggestion_id,
                     organization_id=organization_id,
                 )
                 session.add(new_artifact)
@@ -1789,6 +1793,21 @@ class AgentDB:
             await session.refresh(new_task_generation)
             return TaskGeneration.model_validate(new_task_generation)
 
+    async def create_ai_suggestion(
+        self,
+        organization_id: str,
+        ai_suggestion_type: str,
+    ) -> AISuggestion:
+        async with self.Session() as session:
+            new_ai_suggestion = AISuggestionModel(
+                organization_id=organization_id,
+                ai_suggestion_type=ai_suggestion_type,
+            )
+            session.add(new_ai_suggestion)
+            await session.commit()
+            await session.refresh(new_ai_suggestion)
+            return AISuggestion.model_validate(new_ai_suggestion)
+
     async def get_task_generation_by_prompt_hash(
         self,
         user_prompt_hash: str,

skyvern/forge/sdk/db/id.py

@@ -44,6 +44,7 @@ BITWARDEN_LOGIN_CREDENTIAL_PARAMETER_PREFIX = "blc"
 BITWARDEN_SENSITIVE_INFORMATION_PARAMETER_PREFIX = "bsi"
 BITWARDEN_CREDIT_CARD_DATA_PARAMETER_PREFIX = "bccd"
 TASK_GENERATION_PREFIX = "tg"
+AI_SUGGESTION_PREFIX = "as"
 OBSERVER_CRUISE_ID = "oc"
 OBSERVER_THOUGHT_ID = "ot"
 PERSISTENT_BROWSER_SESSION_ID = "pbs"
@@ -134,6 +135,11 @@ def generate_task_generation_id() -> str:
     return f"{TASK_GENERATION_PREFIX}_{int_id}"
 
 
+def generate_ai_suggestion_id() -> str:
+    int_id = generate_id()
+    return f"{AI_SUGGESTION_PREFIX}_{int_id}"
+
+
 def generate_totp_code_id() -> str:
     int_id = generate_id()
     return f"totp_{int_id}"

skyvern/forge/sdk/db/models.py

@@ -20,6 +20,7 @@ from sqlalchemy.orm import DeclarativeBase
 from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType, TaskType
 from skyvern.forge.sdk.db.id import (
     generate_action_id,
+    generate_ai_suggestion_id,
     generate_artifact_id,
     generate_aws_secret_parameter_id,
     generate_bitwarden_credit_card_data_parameter_id,
@@ -169,6 +170,7 @@ class ArtifactModel(Base):
     workflow_run_block_id = Column(String, index=True)
     observer_cruise_id = Column(String, index=True)
     observer_thought_id = Column(String, index=True)
+    ai_suggestion_id = Column(String, index=True)
     task_id = Column(String, ForeignKey("tasks.task_id"))
     step_id = Column(String, ForeignKey("steps.step_id"), index=True)
     artifact_type = Column(String)
@@ -441,6 +443,16 @@ class TaskGenerationModel(Base):
     modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)
 
 
+class AISuggestionModel(Base):
+    __tablename__ = "ai_suggestions"
+
+    ai_suggestion_id = Column(String, primary_key=True, default=generate_ai_suggestion_id)
+    organization_id = Column(String, ForeignKey("organizations.organization_id"))
+    ai_suggestion_type = Column(String)
+    created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
+    modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)
+
+
 class TOTPCodeModel(Base):
     __tablename__ = "totp_codes"