Update SAML schema to support tools in oss llms (#849)

Author: elisalimli

libs/superagent/app/api/workflow_configs/saml_schema.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Any, Optional
+from typing import Any, Generic, Optional, TypeVar
 
 from pydantic import BaseModel, Field, validator
 
@@ -72,9 +72,9 @@ class ToolModel(BaseModel):
     metaphor: Optional[Tool]
     function: Optional[Tool]
     # ~~~~~~Assistants as tools~~~~~~
-    superagent: Optional["AgentTool"]
-    openai_assistant: Optional["AgentTool"]
-    llm: Optional["AssistantTool"]
+    superagent: Optional["BaseAssistantToolModel[SuperagentAgent]"]
+    openai_assistant: Optional["BaseAssistantToolModel[OpenAIAgent]"]
+    llm: Optional["BaseAssistantToolModel[LLMAgent]"]
 
     # OpenAI Assistant tools
     code_interpreter: Optional[Tool]
@@ -92,22 +92,27 @@ class Assistant(BaseModel):
     intro: Optional[str]
 
 
-class Agent(Assistant):
+# ~~~Agents~~~
+class SuperagentAgent(Assistant):
     tools: Optional[Tools]
-    data: Optional[Data]
+    data: Optional[Data]  # deprecated, use superrag instead
     superrag: Optional[Superrag]
 
 
-class BaseAssistantToolModel(BaseModel):
-    use_for: str
+class LLMAgent(Assistant):
+    tools: Optional[Tools]
+    superrag: Optional[Superrag]
 
 
-class AgentTool(BaseAssistantToolModel, Agent):
+class OpenAIAgent(Assistant):
     pass
 
 
-class AssistantTool(BaseAssistantToolModel, Assistant):
-    pass
+AgentT = TypeVar("AgentT")
+
+
+class BaseAssistantToolModel(BaseModel, Generic[AgentT]):
+    use_for: str
 
 
 # This is for the circular reference between Agent, Assistant and ToolModel
@@ -116,9 +121,9 @@ class AssistantTool(BaseAssistantToolModel, Assistant):
 
 
 class Workflow(BaseModel):
-    superagent: Optional[Agent]
-    openai_assistant: Optional[Assistant]
-    llm: Optional[Assistant]
+    superagent: Optional[SuperagentAgent]
+    openai_assistant: Optional[OpenAIAgent]
+    llm: Optional[LLMAgent]
 
 
 class WorkflowConfigModel(BaseModel):

libs/superagent/app/api/workflow_configs/validator.py

@@ -15,7 +15,7 @@ def validate(self):
         self.validate_superrag_names()
 
     def validate_assistant_names(self):
-        assistants = self.config.get("workflows", [])
+        assistants = self.config.get("workflows") or []
         assistant_names = []
 
         for new_assistant in assistants:
@@ -30,13 +30,13 @@ def validate_assistant_names(self):
             assistant_names.append(assistant_name)
 
     def validate_tool_names(self):
-        assistants = self.config.get("workflows", [])
+        assistants = self.config.get("workflows") or []
 
         for new_assistant in assistants:
             assistant_type = get_first_non_null_key(new_assistant)
             assistant = new_assistant.get(assistant_type)
 
-            tools = assistant.get("tools", [])
+            tools = assistant.get("tools") or []
             tool_names = []
 
             for tool in tools:
@@ -51,7 +51,7 @@ def validate_tool_names(self):
                 tool_names.append(tool_name)
 
     def validate_superrag_names(self):
-        assistants = self.config.get("workflows", [])
+        assistants = self.config.get("workflows") or []
 
         for new_assistant in assistants:
             assistant_type = get_first_non_null_key(new_assistant)