Personal RAG

package main

import (
	"context"
	"crypto/sha256"
	"crypto/subtle"
	"database/sql"
	"errors"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/modelcontextprotocol/go-sdk/mcp"
	_ "modernc.org/sqlite"
)

type App struct {
	db     *sql.DB
	dbPath string
}

type AddTextInput struct {
	Title  string   `json:"title" jsonschema:"short title for the knowledge item"`
	Text   string   `json:"text" jsonschema:"text to store"`
	Source string   `json:"source,omitempty" jsonschema:"optional source, e.g. conversation, note, file"`
	Tags   []string `json:"tags,omitempty" jsonschema:"optional tags"`
}
type AddTextOutput struct {
	ID            int64  `json:"id"`
	AlreadyExists bool   `json:"already_exists"`
	Message       string `json:"message,omitempty"`
}

type AddFactInput struct {
	Subject    string   `json:"subject" jsonschema:"entity or topic the fact is about"`
	Predicate  string   `json:"predicate" jsonschema:"property or relationship"`
	Value      string   `json:"value" jsonschema:"fact value"`
	Source     string   `json:"source,omitempty" jsonschema:"optional source"`
	Confidence float64  `json:"confidence,omitempty" jsonschema:"confidence from 0 to 1; defaults to 1"`
	Tags       []string `json:"tags,omitempty" jsonschema:"optional tags"`
}
type AddFactOutput struct {
	ID            int64  `json:"id"`
	AlreadyExists bool   `json:"already_exists"`
	Message       string `json:"message,omitempty"`
}

type UpdateInput struct {
	ID      int64    `json:"id" jsonschema:"knowledge text item id"`
	Title   *string  `json:"title,omitempty" jsonschema:"new title"`
	Content *string  `json:"content,omitempty" jsonschema:"new content"`
	Tags    []string `json:"tags,omitempty" jsonschema:"new tags"`
}
type UpdateOutput struct {
	Updated bool `json:"updated"`
}

type SearchInput struct {
	Query string `json:"query" jsonschema:"search terms"`
	Limit int    `json:"limit,omitempty" jsonschema:"maximum results, defaults to 5"`
}
type SearchItem struct {
	Kind    string `json:"kind"`
	ID      int64  `json:"id"`
	Title   string `json:"title,omitempty"`
	Content string `json:"content"`
	Source  string `json:"source,omitempty"`
	Tags    string `json:"tags,omitempty"`
}
type SearchOutput struct {
	Results []SearchItem `json:"results"`
}

type DeleteInput struct {
	Kind string `json:"kind" jsonschema:"text or fact"`
	ID   int64  `json:"id" jsonschema:"item id"`
}
type DeleteOutput struct {
	Deleted bool `json:"deleted"`
}

type StatsInput struct{}
type StatsOutput struct {
	TextItems int64  `json:"text_items"`
	Facts     int64  `json:"facts"`
	DBPath    string `json:"db_path"`
}

func main() {
	dbDir := flag.String("db-path", "", "directory containing knowledge.db")
	httpAddr := flag.String("http", "", "serve MCP over Streamable HTTP on this listen address")
	authTokenFlag := flag.String("auth-token", "", "bearer token required for HTTP requests")
	flag.Parse()
	authToken := *authTokenFlag
	if authToken == "" {
		authToken = os.Getenv("PERSONAL_RAG_AUTH_TOKEN")
	}
	if *httpAddr != "" && authToken == "" {
		log.Fatal("HTTP mode requires --auth-token or PERSONAL_RAG_AUTH_TOKEN")
	}
	dbPath, err := resolveDBPath(*dbDir)
	if err != nil {
		log.Fatal(err)
	}
	if err := os.MkdirAll(filepath.Dir(dbPath), 0o755); err != nil {
		log.Fatal(err)
	}

	db, err := sql.Open("sqlite", dbPath)
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	app := &App{db: db, dbPath: dbPath}
	if err := app.migrate(); err != nil {
		log.Fatal(err)
	}

	server := mcp.NewServer(&mcp.Implementation{Name: "personal-rag", Version: "0.1.0"}, nil)
	mcp.AddTool(server, &mcp.Tool{Name: "knowledge_add_text", Description: "Store a text knowledge item. Use only when the user explicitly wants information persisted."}, app.addText)
	mcp.AddTool(server, &mcp.Tool{Name: "knowledge_add_fact", Description: "Store a structured fact as subject/predicate/value. Use only when the user explicitly wants information persisted."}, app.addFact)
	mcp.AddTool(server, &mcp.Tool{Name: "knowledge_search", Description: "Search stored personal knowledge using SQLite FTS5."}, app.search)
	mcp.AddTool(server, &mcp.Tool{Name: "knowledge_update", Description: "Update the title, content, or tags of a stored text item by ID."}, app.update)
	mcp.AddTool(server, &mcp.Tool{Name: "knowledge_delete", Description: "Delete a stored text item or fact by ID."}, app.delete)
	mcp.AddTool(server, &mcp.Tool{Name: "knowledge_stats", Description: "Return counts and database path."}, app.stats)

	if *httpAddr != "" {
		handler := mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server { return server }, nil)
		if err := http.ListenAndServe(*httpAddr, bearerAuth(handler, authToken)); err != nil {
			log.Fatal(err)
		}
	} else {
		if err := server.Run(context.Background(), &mcp.StdioTransport{}); err != nil {
			log.Fatal(err)
		}
	}
}

func bearerAuth(next http.Handler, token string) http.Handler {
	expectedHeader := sha256.Sum256([]byte("Bearer " + token))
	expectedPath := sha256.Sum256([]byte("/token=" + token))
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		providedHeader := sha256.Sum256([]byte(r.Header.Get("Authorization")))
		providedPath := sha256.Sum256([]byte(r.URL.Path))
		headerValid := subtle.ConstantTimeCompare(expectedHeader[:], providedHeader[:])
		pathValid := subtle.ConstantTimeCompare(expectedPath[:], providedPath[:])
		if headerValid|pathValid != 1 {
			http.Error(w, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
			return
		}
		if pathValid == 1 {
			r.URL.Path = "/"
			r.URL.RawPath = ""
		}
		next.ServeHTTP(w, r)
	})
}

func resolveDBPath(dbDir string) (string, error) {
	var path string
	if dbDir != "" {
		path = filepath.Join(dbDir, "knowledge.db")
	} else if legacyPath := os.Getenv("PERSONAL_RAG_DB"); legacyPath != "" {
		path = legacyPath
	} else {
		path = "knowledge.db"
	}
	return filepath.Abs(path)
}

func (a *App) migrate() error {
	schema := `
PRAGMA journal_mode=WAL;
PRAGMA foreign_keys=ON;

CREATE TABLE IF NOT EXISTS knowledge_text (
  id INTEGER PRIMARY KEY,
  title TEXT NOT NULL,
  content TEXT NOT NULL,
  source TEXT NOT NULL DEFAULT '',
  tags TEXT NOT NULL DEFAULT '',
  created_at TEXT NOT NULL,
  updated_at TEXT NOT NULL
);
CREATE VIRTUAL TABLE IF NOT EXISTS knowledge_text_fts USING fts5(
  title, content, source, tags,
  content='knowledge_text', content_rowid='id'
);
CREATE TRIGGER IF NOT EXISTS knowledge_text_ai AFTER INSERT ON knowledge_text BEGIN
  INSERT INTO knowledge_text_fts(rowid,title,content,source,tags)
  VALUES (new.id,new.title,new.content,new.source,new.tags);
END;
CREATE TRIGGER IF NOT EXISTS knowledge_text_ad AFTER DELETE ON knowledge_text BEGIN
  INSERT INTO knowledge_text_fts(knowledge_text_fts,rowid,title,content,source,tags)
  VALUES('delete',old.id,old.title,old.content,old.source,old.tags);
END;
CREATE TRIGGER IF NOT EXISTS knowledge_text_au AFTER UPDATE ON knowledge_text BEGIN
  INSERT INTO knowledge_text_fts(knowledge_text_fts,rowid,title,content,source,tags)
  VALUES('delete',old.id,old.title,old.content,old.source,old.tags);
  INSERT INTO knowledge_text_fts(rowid,title,content,source,tags)
  VALUES (new.id,new.title,new.content,new.source,new.tags);
END;

CREATE TABLE IF NOT EXISTS knowledge_fact (
  id INTEGER PRIMARY KEY,
  subject TEXT NOT NULL,
  predicate TEXT NOT NULL,
  value TEXT NOT NULL,
  source TEXT NOT NULL DEFAULT '',
  confidence REAL NOT NULL DEFAULT 1.0,
  tags TEXT NOT NULL DEFAULT '',
  created_at TEXT NOT NULL,
  updated_at TEXT NOT NULL
);
CREATE VIRTUAL TABLE IF NOT EXISTS knowledge_fact_fts USING fts5(
  subject, predicate, value, source, tags,
  content='knowledge_fact', content_rowid='id'
);
CREATE TRIGGER IF NOT EXISTS knowledge_fact_ai AFTER INSERT ON knowledge_fact BEGIN
  INSERT INTO knowledge_fact_fts(rowid,subject,predicate,value,source,tags)
  VALUES (new.id,new.subject,new.predicate,new.value,new.source,new.tags);
END;
CREATE TRIGGER IF NOT EXISTS knowledge_fact_ad AFTER DELETE ON knowledge_fact BEGIN
  INSERT INTO knowledge_fact_fts(knowledge_fact_fts,rowid,subject,predicate,value,source,tags)
  VALUES('delete',old.id,old.subject,old.predicate,old.value,old.source,old.tags);
END;
CREATE TRIGGER IF NOT EXISTS knowledge_fact_au AFTER UPDATE ON knowledge_fact BEGIN
  INSERT INTO knowledge_fact_fts(knowledge_fact_fts,rowid,subject,predicate,value,source,tags)
  VALUES('delete',old.id,old.subject,old.predicate,old.value,old.source,old.tags);
  INSERT INTO knowledge_fact_fts(rowid,subject,predicate,value,source,tags)
  VALUES (new.id,new.subject,new.predicate,new.value,new.source,new.tags);
END;`
	_, err := a.db.Exec(schema)
	return err
}

func (a *App) addText(ctx context.Context, req *mcp.CallToolRequest, in AddTextInput) (*mcp.CallToolResult, AddTextOutput, error) {
	in.Title = strings.TrimSpace(in.Title)
	in.Text = strings.TrimSpace(in.Text)
	if in.Title == "" || in.Text == "" {
		return nil, AddTextOutput{}, errors.New("title and text are required")
	}
	var existingID int64
	err := a.db.QueryRowContext(ctx, `SELECT id FROM knowledge_text WHERE content=? ORDER BY id LIMIT 1`, in.Text).Scan(&existingID)
	if err == nil {
		return nil, AddTextOutput{ID: existingID, AlreadyExists: true, Message: "entry already exists"}, nil
	}
	if !errors.Is(err, sql.ErrNoRows) {
		return nil, AddTextOutput{}, err
	}
	now := time.Now().UTC().Format(time.RFC3339)
	r, err := a.db.ExecContext(ctx, `INSERT INTO knowledge_text(title,content,source,tags,created_at,updated_at) VALUES(?,?,?,?,?,?)`, in.Title, in.Text, strings.TrimSpace(in.Source), strings.Join(in.Tags, ","), now, now)
	if err != nil {
		return nil, AddTextOutput{}, err
	}
	id, _ := r.LastInsertId()
	return nil, AddTextOutput{ID: id}, nil
}

func (a *App) addFact(ctx context.Context, req *mcp.CallToolRequest, in AddFactInput) (*mcp.CallToolResult, AddFactOutput, error) {
	in.Subject = strings.TrimSpace(in.Subject)
	in.Predicate = strings.TrimSpace(in.Predicate)
	in.Value = strings.TrimSpace(in.Value)
	if in.Subject == "" || in.Predicate == "" || in.Value == "" {
		return nil, AddFactOutput{}, errors.New("subject, predicate and value are required")
	}
	if in.Confidence == 0 {
		in.Confidence = 1
	}
	if in.Confidence < 0 || in.Confidence > 1 {
		return nil, AddFactOutput{}, errors.New("confidence must be between 0 and 1")
	}
	var existingID int64
	err := a.db.QueryRowContext(ctx, `SELECT id FROM knowledge_fact WHERE subject=? AND predicate=? AND value=? ORDER BY id LIMIT 1`, in.Subject, in.Predicate, in.Value).Scan(&existingID)
	if err == nil {
		return nil, AddFactOutput{ID: existingID, AlreadyExists: true, Message: "entry already exists"}, nil
	}
	if !errors.Is(err, sql.ErrNoRows) {
		return nil, AddFactOutput{}, err
	}
	now := time.Now().UTC().Format(time.RFC3339)
	r, err := a.db.ExecContext(ctx, `INSERT INTO knowledge_fact(subject,predicate,value,source,confidence,tags,created_at,updated_at) VALUES(?,?,?,?,?,?,?,?)`, in.Subject, in.Predicate, in.Value, strings.TrimSpace(in.Source), in.Confidence, strings.Join(in.Tags, ","), now, now)
	if err != nil {
		return nil, AddFactOutput{}, err
	}
	id, _ := r.LastInsertId()
	return nil, AddFactOutput{ID: id}, nil
}

func (a *App) update(ctx context.Context, req *mcp.CallToolRequest, in UpdateInput) (*mcp.CallToolResult, UpdateOutput, error) {
	if in.ID <= 0 {
		return nil, UpdateOutput{}, errors.New("id must be positive")
	}
	if in.Title == nil && in.Content == nil && in.Tags == nil {
		return nil, UpdateOutput{}, errors.New("at least one of title, content or tags is required")
	}
	var title, content, tags string
	if err := a.db.QueryRowContext(ctx, `SELECT title,content,tags FROM knowledge_text WHERE id=?`, in.ID).Scan(&title, &content, &tags); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, UpdateOutput{Updated: false}, nil
		}
		return nil, UpdateOutput{}, err
	}
	if in.Title != nil {
		title = strings.TrimSpace(*in.Title)
		if title == "" {
			return nil, UpdateOutput{}, errors.New("title must not be empty")
		}
	}
	if in.Content != nil {
		content = strings.TrimSpace(*in.Content)
		if content == "" {
			return nil, UpdateOutput{}, errors.New("content must not be empty")
		}
	}
	if in.Tags != nil {
		tags = strings.Join(in.Tags, ",")
	}
	_, err := a.db.ExecContext(ctx, `UPDATE knowledge_text SET title=?,content=?,tags=?,updated_at=? WHERE id=?`, title, content, tags, time.Now().UTC().Format(time.RFC3339), in.ID)
	if err != nil {
		return nil, UpdateOutput{}, err
	}
	return nil, UpdateOutput{Updated: true}, nil
}

func ftsQuery(q string) string {
	parts := strings.Fields(q)
	out := make([]string, 0, len(parts))
	for _, p := range parts {
		p = strings.ReplaceAll(p, `"`, "")
		p = strings.TrimSpace(p)
		if p != "" {
			out = append(out, `"`+p+`"`)
		}
	}
	return strings.Join(out, " OR ")
}

func (a *App) search(ctx context.Context, req *mcp.CallToolRequest, in SearchInput) (*mcp.CallToolResult, SearchOutput, error) {
	if strings.TrimSpace(in.Query) == "" {
		return nil, SearchOutput{}, errors.New("query is required")
	}
	if in.Limit <= 0 {
		in.Limit = 5
	}
	if in.Limit > 50 {
		in.Limit = 50
	}
	q := ftsQuery(in.Query)
	results := []SearchItem{}

	rows, err := a.db.QueryContext(ctx, `SELECT k.id,k.title,k.content,k.source,k.tags FROM knowledge_text_fts f JOIN knowledge_text k ON k.id=f.rowid WHERE knowledge_text_fts MATCH ? ORDER BY bm25(knowledge_text_fts) LIMIT ?`, q, in.Limit)
	if err != nil {
		return nil, SearchOutput{}, err
	}
	for rows.Next() {
		var x SearchItem
		x.Kind = "text"
		if err := rows.Scan(&x.ID, &x.Title, &x.Content, &x.Source, &x.Tags); err != nil {
			rows.Close()
			return nil, SearchOutput{}, err
		}
		results = append(results, x)
	}
	rows.Close()

	rows, err = a.db.QueryContext(ctx, `SELECT k.id,k.subject,k.predicate,k.value,k.source,k.tags FROM knowledge_fact_fts f JOIN knowledge_fact k ON k.id=f.rowid WHERE knowledge_fact_fts MATCH ? ORDER BY bm25(knowledge_fact_fts) LIMIT ?`, q, in.Limit)
	if err != nil {
		return nil, SearchOutput{}, err
	}
	for rows.Next() {
		var id int64
		var s, p, v, src, tags string
		if err := rows.Scan(&id, &s, &p, &v, &src, &tags); err != nil {
			rows.Close()
			return nil, SearchOutput{}, err
		}
		results = append(results, SearchItem{Kind: "fact", ID: id, Title: s, Content: fmt.Sprintf("%s: %s", p, v), Source: src, Tags: tags})
	}
	rows.Close()
	if len(results) > in.Limit {
		results = results[:in.Limit]
	}
	return nil, SearchOutput{Results: results}, nil
}

func (a *App) delete(ctx context.Context, req *mcp.CallToolRequest, in DeleteInput) (*mcp.CallToolResult, DeleteOutput, error) {
	table := ""
	switch in.Kind {
	case "text":
		table = "knowledge_text"
	case "fact":
		table = "knowledge_fact"
	default:
		return nil, DeleteOutput{}, errors.New("kind must be text or fact")
	}
	r, err := a.db.ExecContext(ctx, "DELETE FROM "+table+" WHERE id=?", in.ID)
	if err != nil {
		return nil, DeleteOutput{}, err
	}
	n, _ := r.RowsAffected()
	return nil, DeleteOutput{Deleted: n > 0}, nil
}

func (a *App) stats(ctx context.Context, req *mcp.CallToolRequest, in StatsInput) (*mcp.CallToolResult, StatsOutput, error) {
	var t, f int64
	if err := a.db.QueryRowContext(ctx, "SELECT count(*) FROM knowledge_text").Scan(&t); err != nil {
		return nil, StatsOutput{}, err
	}
	if err := a.db.QueryRowContext(ctx, "SELECT count(*) FROM knowledge_fact").Scan(&f); err != nil {
		return nil, StatsOutput{}, err
	}
	return nil, StatsOutput{TextItems: t, Facts: f, DBPath: a.dbPath}, nil
}