diff --git a/README.md b/README.md
index f9e94756b..d05448026 100644
--- a/README.md
+++ b/README.md
@@ -75,8 +75,11 @@ codegen notebook
 
 See [Getting Started](https://docs.codegen.com/introduction/getting-started) for a full tutorial.
 
-```
+```python
 from codegen import Codebase
+
+# For AI agent functionality
+from codegen import Agent
 ```
 
 ## Troubleshooting
diff --git a/src/codegen/agents/README.md b/src/codegen/agents/README.md
index 254ed4bc9..8734ba077 100644
--- a/src/codegen/agents/README.md
+++ b/src/codegen/agents/README.md
@@ -15,7 +15,11 @@ pip install codegen
 ### Basic Example
 
 ```python
-from codegen.agents.agent import Agent
+# You can import Agent directly from the top-level package
+from codegen import Agent
+
+# Or use the full import path
+# from codegen.agents.agent import Agent
 
 # Initialize the Agent with your organization ID and API token
 agent = Agent(