potato_agent/agents/
types.rs

1use crate::agents::error::AgentError;
2use crate::agents::provider::openai::OpenAIChatResponse;
3use potato_prompt::{
4    prompt::{PromptContent, Role},
5    Message,
6};
7
8use pyo3::prelude::*;
9use pyo3::IntoPyObjectExt;
10use serde::{Deserialize, Serialize};
11
12#[pyclass]
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub enum ChatResponse {
15    OpenAI(OpenAIChatResponse),
16}
17
18impl ChatResponse {
19    pub fn is_empty(&self) -> bool {
20        match self {
21            ChatResponse::OpenAI(resp) => resp.choices.is_empty(),
22        }
23    }
24
25    pub fn to_message(&self, role: Role) -> Result<Vec<Message>, AgentError> {
26        match self {
27            ChatResponse::OpenAI(resp) => {
28                let first_choice = resp
29                    .choices
30                    .first()
31                    .ok_or_else(|| AgentError::ClientNoResponseError)?;
32
33                let message =
34                    PromptContent::Str(first_choice.message.content.clone().unwrap_or_default());
35                Ok(vec![Message::from(message, role)])
36            }
37        }
38    }
39
40    pub fn to_python<'py>(&self, py: Python<'py>) -> Result<Bound<'py, PyAny>, AgentError> {
41        match self {
42            ChatResponse::OpenAI(resp) => Ok(resp.clone().into_bound_py_any(py)?),
43        }
44    }
45}
46
47#[pyclass]
48pub struct AgentResponse {
49    #[pyo3(get)]
50    pub id: String,
51    pub response: ChatResponse,
52}
53
54#[pymethods]
55impl AgentResponse {
56    #[getter]
57    pub fn output(&self) -> String {
58        match &self.response {
59            ChatResponse::OpenAI(resp) => resp.choices.first().map_or("".to_string(), |c| {
60                c.message.content.clone().unwrap_or_default()
61            }),
62        }
63    }
64}
65
66impl AgentResponse {
67    pub fn new(id: String, response: ChatResponse) -> Self {
68        Self { id, response }
69    }
70}
potato_agent/agents/types.rs

potato_agent/agents/
types.rs
