package schema

import (
    "entgo.io/ent"
    "entgo.io/ent/schema/edge"
    "entgo.io/ent/schema/field"
)

// Chunk holds the schema definition for the Chunk entity.
type Chunk struct {
    ent.Schema
}

// Fields of the Chunk.
func (Chunk) Fields() []ent.Field {
    return []ent.Field{
        field.String("path"),
        field.Int("nchunk"),
        field.Text("data"),
    }
}

// Edges of the Chunk.
func (Chunk) Edges() []ent.Edge {
    return []ent.Edge{
        edge.To("embedding", Embedding.Type).StorageKey(edge.Column("chunk_id")).Unique(),
    }
}

package schema

import (
    "entgo.io/ent"
    "entgo.io/ent/dialect"
    "entgo.io/ent/dialect/entsql"
    "entgo.io/ent/schema/edge"
    "entgo.io/ent/schema/field"
    "entgo.io/ent/schema/index"
    "github.com/pgvector/pgvector-go"
)

// Embedding holds the schema definition for the Embedding entity.
type Embedding struct {
    ent.Schema
}

// Fields of the Embedding.
func (Embedding) Fields() []ent.Field {
    return []ent.Field{
        field.Other("embedding", pgvector.Vector{}).
            SchemaType(map[string]string{
                dialect.Postgres: "vector(1536)",
            }),
    }
}

// Edges of the Embedding.
func (Embedding) Edges() []ent.Edge {
    return []ent.Edge{
        edge.From("chunk", Chunk.Type).Ref("embedding").Unique().Required(),
    }
}

func (Embedding) Indexes() []ent.Index {
    return []ent.Index{
        index.Fields("embedding").
            Annotations(
                entsql.IndexType("hnsw"),
                entsql.OpClass("vector_l2_ops"),
            ),
    }
}

package main

import (
    "fmt"
    "os"

    "github.com/alecthomas/kong"
)

// CLI holds global options and subcommands.
type CLI struct {
    // DBURL is read from the environment variable DB_URL.
    DBURL     string `kong:"env='DB_URL',help='Database URL for the application.'"`
    OpenAIKey string `kong:"env='OPENAI_KEY',help='OpenAI API key for the application.'"`

    // Subcommands
    Load  *LoadCmd  `kong:"cmd,help='Load command that accepts a path.'"`
    Index *IndexCmd `kong:"cmd,help='Create embeddings for any chunks that do not have one.'"`
    Ask   *AskCmd   `kong:"cmd,help='Ask a question about the indexed documents'"`
}

func main() {
    var cli CLI
    app := kong.Parse(&cli,
        kong.Name("entrag"),
        kong.Description("Ask questions about markdown files."),
        kong.UsageOnError(),
    )
    if err := app.Run(&cli); err != nil {
        fmt.Fprintf(os.Stderr, "Error: %s\n", err)
        os.Exit(1)
    }
}

package main

type (
    // LoadCmd loads the markdown files into the database.
    LoadCmd struct {
        Path string `help:"path to dir with markdown files" type:"existingdir" required:""`
    }
    // IndexCmd creates the embedding index on the database.
    IndexCmd struct {
    }
    // AskCmd is another leaf command.
    AskCmd struct {
        // Text is the positional argument for the ask command.
        Text string `kong:"arg,required,help='Text for the ask command.'"`
    }
)

const (
    tokenEncoding = "cl100k_base"
    chunkSize     = 1000
)

// Run is the method called when the "load" command is executed.
func (cmd *LoadCmd) Run(ctx *CLI) error {
    client, err := ctx.entClient()
    if err != nil {
        return fmt.Errorf("failed opening connection to postgres: %w", err)
    }
    tokTotal := 0
    return filepath.WalkDir(ctx.Load.Path, func(path string, d fs.DirEntry, err error) error {
        if filepath.Ext(path) == ".mdx" || filepath.Ext(path) == ".md" {
            chunks := breakToChunks(path)
            for i, chunk := range chunks {
                tokTotal += len(chunk)
                client.Chunk.Create().
                    SetData(chunk).
                    SetPath(path).
                    SetNchunk(i).
                    SaveX(context.Background())
            }
        }
        return nil
    })
}

func (c *CLI) entClient() (*ent.Client, error) {
    return ent.Open("postgres", c.DBURL)
}

// Run is the method called when the "index" command is executed.
func (cmd *IndexCmd) Run(cli *CLI) error {
    client, err := cli.entClient()
    if err != nil {
        return fmt.Errorf("failed opening connection to postgres: %w", err)
    }
    ctx := context.Background()
    chunks := client.Chunk.Query().
        Where(
            chunk.Not(
                chunk.HasEmbedding(),
            ),
        ).
        Order(ent.Asc(chunk.FieldID)).
        AllX(ctx)
    for _, ch := range chunks {
        log.Println("Created embedding for chunk", ch.Path, ch.Nchunk)
        embedding := getEmbedding(ch.Data)
        _, err := client.Embedding.Create().
            SetEmbedding(pgvector.NewVector(embedding)).
            SetChunk(ch).
            Save(ctx)
        if err != nil {
            return fmt.Errorf("error creating embedding: %v", err)
        }
    }
    return nil
}

// getEmbedding invokes the OpenAI embedding API to calculate the embedding
// for the given string. It returns the embedding.
func getEmbedding(data string) []float32 {
    client := openai.NewClient(os.Getenv("OPENAI_KEY"))
    queryReq := openai.EmbeddingRequest{
        Input: []string{data},
        Model: openai.AdaEmbeddingV2,
    }
    queryResponse, err := client.CreateEmbeddings(context.Background(), queryReq)
    if err != nil {
        log.Fatalf("Error getting embedding: %v", err)
    }
    return queryResponse.Data[0].Embedding
}

// Run is the method called when the "ask" command is executed.
func (cmd *AskCmd) Run(ctx *CLI) error {
    client, err := ctx.entClient()
    if err != nil {
        return fmt.Errorf("failed opening connection to postgres: %w", err)
    }
    question := cmd.Text
    emb := getEmbedding(question)
    embVec := pgvector.NewVector(emb)
    embs := client.Embedding.
        Query().
        Order(func(s *sql.Selector) {
            s.OrderExpr(sql.ExprP("embedding <-> $1", embVec))
        }).
        WithChunk().
        Limit(5).
        AllX(context.Background())
    b := strings.Builder{}
    for _, e := range embs {
        chnk := e.Edges.Chunk
        b.WriteString(fmt.Sprintf("From file: %v\n", chnk.Path))
        b.WriteString(chnk.Data)
    }
    query := fmt.Sprintf(`Use the below information from the ent docs to answer the subsequent question.
Information:
%v

Question: %v`, b.String(), question)
    oac := openai.NewClient(ctx.OpenAIKey)
    resp, err := oac.CreateChatCompletion(
        context.Background(),
        openai.ChatCompletionRequest{
            Model: openai.GPT4o,
            Messages: []openai.ChatCompletionMessage{

                {
                    Role:    openai.ChatMessageRoleUser,
                    Content: query,
                },
            },
        },
    )
    if err != nil {
        return fmt.Errorf("error creating chat completion: %v", err)
    }
    choice := resp.Choices[0]
    out, err := glamour.Render(choice.Message.Content, "dark")
    fmt.Print(out)
    return nil
}

emb := getEmbedding(question)
embVec := pgvector.NewVector(emb)
embs := client.Embedding.
    Query().
    Order(func(s *sql.Selector) {
        s.OrderExpr(sql.ExprP("embedding <-> $1", embVec))
    }).
    WithChunk().
    Limit(5).
    AllX(context.Background())

for _, e := range embs {
        chnk := e.Edges.Chunk
        b.WriteString(fmt.Sprintf("From file: %v\n", chnk.Path))
        b.WriteString(chnk.Data)
    }
    query := fmt.Sprintf(`Use the below information from the ent docs to answer the subsequent question.
Information:
%v

Question: %v`, b.String(), question)

oac := openai.NewClient(ctx.OpenAIKey)
resp, err := oac.CreateChatCompletion(
    context.Background(),
    openai.ChatCompletionRequest{
        Model: openai.GPT4o,
        Messages: []openai.ChatCompletionMessage{

            {
                Role:    openai.ChatMessageRoleUser,
                Content: query,
            },
        },
    },
)
if err != nil {
    return fmt.Errorf("error creating chat completion: %v", err)
}
choice := resp.Choices[0]
out, err := glamour.Render(choice.Message.Content, "dark")
fmt.Print(out)

// Fields of the User.  
func (User) Fields() []ent.Field {  
   return []ent.Field{  
      field.String("name"),  
      field.String("email").
            Unique(),  
      field.Time("created_at").
            Default(time.Now),  
   }  
}  
  
// Edges of the User.  
func (User) Edges() []ent.Edge {  
   return []ent.Edge{  
      edge.To("posts", Post.Type),  
   }  
}

// Post holds the schema definition for the Post entity.
type Post struct {  
   ent.Schema  
}  
  
// Fields of the Post.  
func (Post) Fields() []ent.Field {  
   return []ent.Field{  
      field.String("title"),  
      field.Text("body"),  
      field.Time("created_at").
            Default(time.Now),  
   }  
}  
  
// Edges of the Post.
func (Post) Edges() []ent.Edge {  
   return []ent.Edge{  
      edge.From("author", User.Type).  
         Unique().  
         Ref("posts"),  
   }  
}

package main

import (
    "context"
    "flag"
    "fmt"
    "log"

    "github.com/rotemtam/ent-blog-example/ent"

    _ "github.com/go-sql-driver/mysql"
    "github.com/rotemtam/ent-blog-example/ent/user"
)

func main() {
    // Read the connection string to the database from a CLI flag.
    var dsn string
    flag.StringVar(&dsn, "dsn", "", "database DSN")
    flag.Parse()

    // Instantiate the Ent client.
    client, err := ent.Open("mysql", dsn)
    if err != nil {
        log.Fatalf("failed connecting to mysql: %v", err)
    }
    defer client.Close()

    ctx := context.Background()
    // If we don't have any posts yet, seed the database.
    if !client.Post.Query().ExistX(ctx) {
        if err := seed(ctx, client); err != nil {
            log.Fatalf("failed seeding the database: %v", err)
        }
    }
    // ... Continue with server start.
}

func seed(ctx context.Context, client *ent.Client) error {
    // Check if the user "rotemtam" already exists.
    r, err := client.User.Query().
        Where(
            user.Name("rotemtam"),
        ).
        Only(ctx)
    switch {
    // If not, create the user.
    case ent.IsNotFound(err):
        r, err = client.User.Create().
            SetName("rotemtam").
            SetEmail("r@hello.world").
            Save(ctx)
        if err != nil {
            return fmt.Errorf("failed creating user: %v", err)
        }
    case err != nil:
        return fmt.Errorf("failed querying user: %v", err)
    }
    // Finally, create a "Hello, world" blogpost.
    return client.Post.Create().
        SetTitle("Hello, World!").
        SetBody("This is my first post").
        SetAuthor(r).
        Exec(ctx)
}

var (  
   //go:embed templates/*   
   resources embed.FS  
   tmpl      = template.Must(template.ParseFS(resources, "templates/*"))  
)

type server struct {
    client *ent.Client
}

func newServer(client *ent.Client) *server {
    return &server{client: client}
}

// index serves the blog home page
func (s *server) index(w http.ResponseWriter, r *http.Request) {
    posts, err := s.client.Post.
        Query().
        WithAuthor().
        All(r.Context())
    if err != nil {
        http.Error(w, err.Error(), http.StatusInternalServerError)
        return
    }
    if err := tmpl.Execute(w, posts); err != nil {
        http.Error(w, err.Error(), http.StatusInternalServerError)
    }
}

// s.client.Post contains methods for interacting with Post entities
// s.client.Post 包含与 Post 实体交互的方法
s.client.Post.
    // Begin a query.
    // 开始查询。
    Query().
    // Retrieve the entities using the `Author` edge. (a `User` instance)
    // 使用 `Author` 边检索实体。（一个 `User` 实例）
    WithAuthor().
    // Run the query against the database using the request context.
    // 使用请求上下文对数据库运行查询。
    All(r.Context())

// newRouter creates a new router with the blog handlers mounted.
func newRouter(srv *server) chi.Router {
    r := chi.NewRouter()
    r.Use(middleware.Logger)
    r.Use(middleware.Recoverer)
    r.Get("/", srv.index)
    return r
}

package main

import (
    "context"
    "io"
    "net/http"
    "net/http/httptest"
    "testing"

    _ "github.com/mattn/go-sqlite3"
    "github.com/rotemtam/ent-blog-example/ent/enttest"
    "github.com/stretchr/testify/require"
)

func TestIndex(t *testing.T) {
    // Initialize an Ent client that uses an in memory SQLite db.
    client := enttest.Open(t, "sqlite3", "file:ent?mode=memory&cache=shared&_fk=1")
    defer client.Close()
    
    // seed the database with our "Hello, world" post and user.
    err := seed(context.Background(), client)
    require.NoError(t, err)

    // Initialize a server and router.
    srv := newServer(client)
    r := newRouter(srv)

    // Create a test server using the `httptest` package.
    ts := httptest.NewServer(r)
    defer ts.Close()

    // Make a GET request to the server root path.
    resp, err := ts.Client().Get(ts.URL)

    // Assert we get a 200 OK status code.
    require.NoError(t, err)
    require.Equal(t, http.StatusOK, resp.StatusCode)

    // Read the response body and assert it contains "Hello, world!"
    body, err := io.ReadAll(resp.Body)
    require.NoError(t, err)
    require.Contains(t, string(body), "Hello, World!")
}

func main() {  
   // Read the connection string to the database from a CLI flag.  
   var dsn string  
   flag.StringVar(&dsn, "dsn", "", "database DSN")  
   flag.Parse()  
  
   // Instantiate the Ent client.  
   client, err := ent.Open("mysql", dsn)  
   if err != nil {  
      log.Fatalf("failed connecting to mysql: %v", err)  
   }  
   defer client.Close()  
  
   ctx := context.Background()  
   // If we don't have any posts yet, seed the database.  
   if !client.Post.Query().ExistX(ctx) {  
      if err := seed(ctx, client); err != nil {  
         log.Fatalf("failed seeding the database: %v", err)  
      }  
   }  
   srv := newServer(client)  
   r := newRouter(srv)  
   log.Fatal(http.ListenAndServe(":8080", r))  
}

// add creates a new blog post.
func (s *server) add(w http.ResponseWriter, r *http.Request) {
    author, err := s.client.User.Query().Only(r.Context())
    if err != nil {
        http.Error(w, err.Error(), http.StatusInternalServerError)
        return
    }
    if err := s.client.Post.Create().
        SetTitle(r.FormValue("title")).
        SetBody(r.FormValue("body")).
        SetAuthor(author).
        Exec(r.Context()); err != nil {
        http.Error(w, err.Error(), http.StatusInternalServerError)
    }
    http.Redirect(w, r, "/", http.StatusFound)
}

// newRouter creates a new router with the blog handlers mounted.
func newRouter(srv *server) chi.Router {
    r := chi.NewRouter()
    r.Use(middleware.Logger)
    r.Use(middleware.Recoverer)
    r.Get("/", srv.index)
    r.Post("/add", srv.add)
    return r
}