diff --git a/api/graph.go b/api/graph.go
index 15213fef..fcf39d29 100644
--- a/api/graph.go
+++ b/api/graph.go
@@ -3,21 +3,21 @@ package api
 import (
 	"bytes"
 	"fmt"
+	"github.com/gin-gonic/gin"
+	"github.com/smira/aptly/deb"
 	"io"
 	"os"
 	"os/exec"
-	"github.com/gin-gonic/gin"
-	"github.com/smira/aptly/graph"
 )
 
 // GET /api/graph
 func apiGraph(c *gin.Context) {
 	var (
-		err      error
+		err    error
 		output []byte
 	)
 
-	graph, err := graph.BuildGraph(context)
+	graph, err := deb.BuildGraph(context.CollectionFactory())
 	if err != nil {
 		c.JSON(500, err)
 		return
diff --git a/cmd/graph.go b/cmd/graph.go
index 91f2e1a0..552066da 100644
--- a/cmd/graph.go
+++ b/cmd/graph.go
@@ -3,8 +3,8 @@ package cmd
 import (
 	"bytes"
 	"fmt"
+	"github.com/smira/aptly/deb"
 	"github.com/smira/commander"
-	"github.com/smira/aptly/graph"
 	"io"
 	"io/ioutil"
 	"os"
@@ -20,7 +20,7 @@ func aptlyGraph(cmd *commander.Command, args []string) error {
 	}
 
 	fmt.Printf("Generating graph...\n")
-	graph, err := graph.BuildGraph(context)
+	graph, err := deb.BuildGraph(context.CollectionFactory())
 	if err != nil {
 		return err
 	}
diff --git a/graph/graph.go b/deb/graph.go
similarity index 72%
rename from graph/graph.go
rename to deb/graph.go
index 7f9e9daf..108b6d95 100644
--- a/graph/graph.go
+++ b/deb/graph.go
@@ -1,14 +1,12 @@
-package graph
+package deb
 
 import (
 	"code.google.com/p/gographviz"
 	"fmt"
 	"strings"
-	"github.com/smira/aptly/context"
-	"github.com/smira/aptly/deb"
 )
 
-func BuildGraph(context *context.AptlyContext) (gographviz.Interface, error) {
+func BuildGraph(collectionFactory *CollectionFactory) (gographviz.Interface, error) {
 	var err error
 
 	graph := gographviz.NewEscape()
@@ -17,8 +15,8 @@ func BuildGraph(context *context.AptlyContext) (gographviz.Interface, error) {
 
 	existingNodes := map[string]bool{}
 
-	err = context.CollectionFactory().RemoteRepoCollection().ForEach(func(repo *deb.RemoteRepo) error {
-		err := context.CollectionFactory().RemoteRepoCollection().LoadComplete(repo)
+	err = collectionFactory.RemoteRepoCollection().ForEach(func(repo *RemoteRepo) error {
+		err := collectionFactory.RemoteRepoCollection().LoadComplete(repo)
 		if err != nil {
 			return err
 		}
@@ -39,8 +37,8 @@ func BuildGraph(context *context.AptlyContext) (gographviz.Interface, error) {
 		return nil, err
 	}
 
-	err = context.CollectionFactory().LocalRepoCollection().ForEach(func(repo *deb.LocalRepo) error {
-		err := context.CollectionFactory().LocalRepoCollection().LoadComplete(repo)
+	err = collectionFactory.LocalRepoCollection().ForEach(func(repo *LocalRepo) error {
+		err := collectionFactory.LocalRepoCollection().LoadComplete(repo)
 		if err != nil {
 			return err
 		}
@@ -60,13 +58,13 @@ func BuildGraph(context *context.AptlyContext) (gographviz.Interface, error) {
 		return nil, err
 	}
 
-	context.CollectionFactory().SnapshotCollection().ForEach(func(snapshot *deb.Snapshot) error {
+	collectionFactory.SnapshotCollection().ForEach(func(snapshot *Snapshot) error {
 		existingNodes[snapshot.UUID] = true
 		return nil
 	})
 
-	err = context.CollectionFactory().SnapshotCollection().ForEach(func(snapshot *deb.Snapshot) error {
-		err := context.CollectionFactory().SnapshotCollection().LoadComplete(snapshot)
+	err = collectionFactory.SnapshotCollection().ForEach(func(snapshot *Snapshot) error {
+		err := collectionFactory.SnapshotCollection().LoadComplete(snapshot)
 		if err != nil {
 			return err
 		}
@@ -98,7 +96,7 @@ func BuildGraph(context *context.AptlyContext) (gographviz.Interface, error) {
 		return nil, err
 	}
 
-	context.CollectionFactory().PublishedRepoCollection().ForEach(func(repo *deb.PublishedRepo) error {
+	collectionFactory.PublishedRepoCollection().ForEach(func(repo *PublishedRepo) error {
 		graph.AddNode("aptly", repo.UUID, map[string]string{
 			"shape":     "Mrecord",
 			"style":     "filled",
@@ -118,4 +116,4 @@ func BuildGraph(context *context.AptlyContext) (gographviz.Interface, error) {
 	})
 
 	return graph, nil
-}
\ No newline at end of file
+}