diff --git a/README.md b/README.md
index 4a011e5..ab5b4f6 100644
--- a/README.md
+++ b/README.md
@@ -9,15 +9,14 @@
 This is a simple example of how to use the library to create a simple chatbot that uses OpenAI to generate responses.
 
 ```go
+package main
+
 import (
 	"context"
-	"errors"
 	"fmt"
-	"testing"
 
 	"github.com/tmc/langchaingo/llms"
 	"github.com/tmc/langchaingo/llms/openai"
-	"github.com/tmc/langchaingo/schema"
 	"github.com/tmc/langgraphgo/graph"
 )
 
@@ -35,7 +34,7 @@ func main() {
 			return nil, err
 		}
 		return append(state,
-			llms.TextParts(schema.ChatMessageTypeAI, r.Choices[0].Content),
+			llms.TextParts(llms.ChatMessageTypeAI, r.Choices[0].Content),
 		), nil
 
 	})
@@ -54,7 +53,7 @@ func main() {
 	ctx := context.Background()
 	// Let's run it!
 	res, err := runnable.Invoke(ctx, []llms.MessageContent{
-		llms.TextParts(schema.ChatMessageTypeHuman, "What is 1 + 1?"),
+		llms.TextParts(llms.ChatMessageTypeHuman, "What is 1 + 1?"),
 	})
 	if err != nil {
 		panic(err)
@@ -63,6 +62,6 @@ func main() {
 	fmt.Println(res)
 
 	// Output:
-	// [{human [{What is 1 + 1?}]} {ai [{1 + 1 equals 2.}]}]
+	// [{human [What is 1 + 1?]} {ai [The answer to 1 + 1 is... 2!]}]
 }
 ```