diff --git a/pkg/client/client.go b/pkg/client/client.go
index 0b750aa..508f7fd 100644
--- a/pkg/client/client.go
+++ b/pkg/client/client.go
@@ -17,12 +17,15 @@ type Client struct {
 }
 
 // NewClient creates a new LocalAgent client
-func NewClient(baseURL string, apiKey string) *Client {
+func NewClient(baseURL string, apiKey string, timeout time.Duration) *Client {
+	if timeout == 0 {
+		timeout = time.Second * 30
+	}
 	return &Client{
 		BaseURL: baseURL,
 		APIKey:  apiKey,
 		HTTPClient: &http.Client{
-			Timeout: time.Second * 30,
+			Timeout: timeout,
 		},
 	}
 }
diff --git a/pkg/client/responses.go b/pkg/client/responses.go
index bf05063..32bb6d2 100644
--- a/pkg/client/responses.go
+++ b/pkg/client/responses.go
@@ -8,17 +8,16 @@ import (
 
 // RequestBody represents the message request to the AI model
 type RequestBody struct {
-	Model         string         `json:"model"`
-	Input         string         `json:"input"`
-	InputMessages []InputMessage `json:"input_messages,omitempty"`
-	Temperature   *float64       `json:"temperature,omitempty"`
-	MaxTokens     *int           `json:"max_output_tokens,omitempty"`
+	Model       string   `json:"model"`
+	Input       any      `json:"input"`
+	Temperature *float64 `json:"temperature,omitempty"`
+	MaxTokens   *int     `json:"max_output_tokens,omitempty"`
 }
 
 // InputMessage represents a user input message
 type InputMessage struct {
-	Role    string        `json:"role"`
-	Content []ContentItem `json:"content"`
+	Role    string `json:"role"`
+	Content any    `json:"content"`
 }
 
 // ContentItem represents an item in a content array
@@ -32,7 +31,7 @@ type ContentItem struct {
 type ResponseBody struct {
 	CreatedAt int64             `json:"created_at"`
 	Status    string            `json:"status"`
-	Error     interface{}       `json:"error,omitempty"`
+	Error     any               `json:"error,omitempty"`
 	Output    []ResponseMessage `json:"output"`
 }
 
@@ -103,9 +102,9 @@ func (c *Client) SimpleAIResponse(agentName, input string) (string, error) {
 func (c *Client) ChatAIResponse(agentName string, messages []InputMessage) (string, error) {
 	temperature := 0.7
 	request := &RequestBody{
-		Model:         agentName,
-		InputMessages: messages,
-		Temperature:   &temperature,
+		Model:       agentName,
+		Input:       messages,
+		Temperature: &temperature,
 	}
 
 	response, err := c.GetAIResponse(request)
diff --git a/tests/e2e/e2e_test.go b/tests/e2e/e2e_test.go
index 9c2b91c..a3518d4 100644
--- a/tests/e2e/e2e_test.go
+++ b/tests/e2e/e2e_test.go
@@ -1,6 +1,8 @@
 package e2e_test
 
 import (
+	"time"
+
 	localagent "github.com/mudler/LocalAgent/pkg/client"
 
 	. "github.com/onsi/ginkgo/v2"
@@ -8,9 +10,9 @@ import (
 )
 
 var _ = Describe("Agent test", func() {
-	Context("Creates an agent and it answer", func() {
+	Context("Creates an agent and it answers", func() {
 		It("create agent", func() {
-			client := localagent.NewClient(localagentURL, "")
+			client := localagent.NewClient(localagentURL, "", 5*time.Minute)
 
 			err := client.CreateAgent(&localagent.AgentConfig{
 				Name: "testagent",