diff --git a/README.md b/README.md
index b270105a0..456ada1cd 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,6 @@
 # Codegen
 
 [![Documentation](https://img.shields.io/badge/docs-docs.codegen.com-blue)](https://docs.codegen.com)
-[![Unit Tests](https://github.com/codegen-sh/codegen-sdk/actions/workflows/unit-tests.yml/badge.svg)](https://github.com/codegen-sh/codegen-sdk/actions/workflows/unit-tests.yml)
 
 [Codegen](https://docs.codegen.com) is a python library for manipulating codebases.
 
@@ -42,7 +41,7 @@ We built Codegen backwards from real-world refactors performed on enterprise cod
 
 - **Natural mental model**: Write transforms that read like your thought process - "move this function", "rename this variable", "add this parameter". No more wrestling with ASTs or manual import management.
 
-- **Battle-tested on complex codebases**: Handle Python, TypeScript, and React codebases with millions of lines of code. Built and validated on refactors at companies like [Ramp](https://ramp.com).
+- **Battle-tested on complex codebases**: Handle Python, TypeScript, and React codebases with millions of lines of code.
 
 - **Built for advanced intelligences**: As AI developers become more sophisticated, they need expressive yet precise tools to manipulate code. Codegen provides a programmatic interface that both humans and AI can use to express complex transformations through code itself.
 
diff --git a/src/codegen/__init__.py b/src/codegen/__init__.py
new file mode 100644
index 000000000..c61714134
--- /dev/null
+++ b/src/codegen/__init__.py
@@ -0,0 +1,5 @@
+from codegen.cli.sdk.decorator import function
+from codegen.cli.sdk.functions import Function
+from codegen.sdk.core.codebase import Codebase
+
+__all__ = ["Codebase", "Function", "function"]
diff --git a/tests/unit/test_imports.py b/tests/unit/test_imports.py
new file mode 100644
index 000000000..e2af00f64
--- /dev/null
+++ b/tests/unit/test_imports.py
@@ -0,0 +1,16 @@
+import codegen
+from codegen import Codebase
+
+
+def test_codegen_imports():
+    # Test decorated function
+    @codegen.function(name="sample_codemod")
+    def run(codebase):
+        pass
+
+    # Test class
+    cls = codegen.Function
+    assert cls is not None
+
+    codebase = Codebase("./")
+    assert codebase is not None