Documentation
¶
Overview ¶
Package rag provides functions to integrate Retrieval Augmented Generation (RAG) capabilities in PostgreSQL tables using the pgvector extension. It facilitates operations for adding embeddings on tables from popular language model APIs such as OpenAI, Ollama, and LMStudio.
Index ¶
- type Client
- func (c *Client) CreateEmbedding(ctx context.Context, contentSelectQuery ...string) error
- func (c *Client) FetchEmbedding(ctx context.Context, input []string) ([][]float32, error)
- func (c *Client) Generate(ctx context.Context, prompt string) ([]byte, error)
- func (c *Client) GenerateWithRetrieval(ctx context.Context, prompt string, retrievalLimit int, ...) ([]byte, error)
- func (c *Client) Retrieve(ctx context.Context, input string, limit int) ([]Embedding, error)
- type Config
- type Embedding
- type EmbeddingRequest
- type EmbeddingResponse
- type GenerateRequest
Constants ¶
This section is empty.
Variables ¶
This section is empty.
Functions ¶
This section is empty.
Types ¶
type Client ¶
type Client struct {
Config Config
// contains filtered or unexported fields
}
Client handles the RAG operations
func (*Client) CreateEmbedding ¶
CreateEmbedding inserts embeddings into the table based on the (optional) contentSelectQuery argument If no contentSelectQuery is supplied, it assumes the content column of the table is already populated, and use content column directly If an empty(`""`) contentSelectQuery is provided, it queries the table rows and constructs the content column as col1_name:col1_value,col2_name:col2_value,... For a non-empty contentSelectQuery, it expects the query to return two columns: primary key and content.
Exectuting contentSelectQuery returns output in below format: primary_key, content; pk can be named anything and of any type, allowing usage of existing indexes on the primary key
postgres=# SELECT id, CONCAT('title:', title, ', summary:', summary) AS content FROM lms.courses;
id | content
--------------+-----------------------------------------------------------------------------------------------------------------
15670e34129c | title:Data Science Fundamentals, summary:Explore statistical analysis, machine learning, and data visualization 69c26a5ace74 | title:Web Development Bootcamp, summary:Learn HTML, CSS, JavaScript, and popular frameworks 7a354de48450 | title:Introduction to Python, summary:A comprehensive course for beginners to start their Python journey
func (*Client) FetchEmbedding ¶
FetchEmbedding fetches embeddings from the LLM API
func (*Client) GenerateWithRetrieval ¶
func (c *Client) GenerateWithRetrieval(ctx context.Context, prompt string, retrievalLimit int, retrievalInput ...string) ([]byte, error)
GenerateWithRetrieval performs retrieval-augmented generation (RAG). It retrieves relevant information based on the given prompt or an optional retrieval query, then uses this information to augment the original prompt before generating a response.
Parameters:
- ctx: The context for the operation, which can be used for cancellation.
- prompt: The main prompt or question to be answered by the language model.
- retrievalLimit: The maximum number of relevant documents to retrieve.
- retrievalInput: An optional query used specifically for retrieving relevant documents. If not provided or empty, the prompt will be used for retrieval.
Returns:
- []byte: The generated response from the language model.
- error: An error if any step in the process fails.
The function follows these steps:
- Retrieve relevant information using either the retrievalInput (if provided) or the prompt.
- Construct an augmented prompt that includes the retrieved information and the original prompt.
- Generate a response using the augmented prompt.
This method allows for more flexible and potentially more accurate responses by incorporating relevant context into the generation process.
type Config ¶
type Config struct {
TableName string
TablePrimaryKeyCol string
ModelID string
APIURL string
APIKey string
EmbeddingsPath string
GeneratePath string
Dimensions int
BatchSize int
}
Config holds the configuration for the RAG Client
type Embedding ¶
type Embedding struct {
// PK is the primary key column name
// pk can be named anything (e.g. id, uuid, etc.) and of any type (e.g. int, string, etc.)
// allowing usage of existing indexes on the primary key
PK interface{}
// Content is what vector embeddings are created from
Content string
// Embedding is the vector embedding for the content
Embedding pgvector.Vector
}
Embedding represents a piece of content (e.g. a document, a sentence, base64-encoded image, etc.) with its embedding and a unique ID (primary key) In a PostgreSQL table, it represents a row with columns: primary key, content, and embedding.
type EmbeddingRequest ¶
EmbeddingRequest is the request body for the FetchEmbedding function
type EmbeddingResponse ¶
type EmbeddingResponse struct {
Data []struct {
Embedding []float32 `json:"embedding"`
} `json:"data"`
}
EmbeddingResponse is the response body for the FetchEmbedding function https://platform.openai.com/docs/api-reference/embeddings/create https://github.com/ollama/ollama/blob/main/docs/api.md#embeddings
type GenerateRequest ¶
type GenerateRequest struct {
KeepAlive *time.Duration `json:"keep_alive,omitempty"`
Options map[string]interface{} `json:"options"`
Model string `json:"model"`
Prompt string `json:"prompt"`
Suffix string `json:"suffix"`
System string `json:"system"`
Template string `json:"template"`
Format string `json:"format"`
Context []int `json:"context,omitempty"`
Images []string `json:"images,omitempty"`
Stream bool `json:"stream"`
Raw bool `json:"raw,omitempty"`
}
GenerateRequest is the body for /generate requests. Model and Prompt fields are required.
Source Files