Initial commit

api/core/callback_handler/entity/agent_loop.py

@@ -0,0 +1,23 @@
+from pydantic import BaseModel
+
+
+class AgentLoop(BaseModel):
+    position: int = 1
+
+    thought: str = None
+    tool_name: str = None
+    tool_input: str = None
+    tool_output: str = None
+
+    prompt: str = None
+    prompt_tokens: int = None
+    completion: str = None
+    completion_tokens: int = None
+
+    latency: float = None
+
+    status: str = 'llm_started'
+    completed: bool = False
+
+    started_at: float = None
+    completed_at: float = None

api/core/callback_handler/entity/chain_result.py

@@ -0,0 +1,16 @@
+from pydantic import BaseModel
+
+
+class ChainResult(BaseModel):
+    type: str = None
+    prompt: dict = None
+    completion: dict = None
+
+    status: str = 'chain_started'
+    completed: bool = False
+
+    started_at: float = None
+    completed_at: float = None
+
+    agent_result: dict = None
+    """only when type is 'AgentExecutor'"""

api/core/callback_handler/entity/dataset_query.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+
+class DatasetQueryObj(BaseModel):
+    dataset_id: str = None
+    query: str = None

api/core/callback_handler/entity/llm_message.py

@@ -0,0 +1,9 @@
+from pydantic import BaseModel
+
+
+class LLMMessage(BaseModel):
+    prompt: str = ''
+    prompt_tokens: int = 0
+    completion: str = ''
+    completion_tokens: int = 0
+    latency: float = 0.0