Merge branch 'release/v1.1.7'

Author: bububa

README.md

@@ -47,7 +47,15 @@ go get -u github.com/bububa/atomic-agents
 Atomic Agents with the following main components:
 
 1. `agents/`: The core Atomic Agents library
-2. `components/`: The Atomic Agents components contains `Message`, `Memory`, `SystemPromptGenerator`, `SystemPromptContextProvider` utilities
+2. `components/`: The Atomic Agents components
+
+- `message`: Defines the Message structure for input/output
+- `memory`: Defines a in memory Memory Store
+- `systemprompt`: Contains SystemPrompt `Generator` and `ContextProvider`
+- `embedder`: Defines the embedder interface, contains several `Provider` including `OpenAI`, `Gemini`, `VoyageAI`, `HuggingFace`, `Cohere` implementations
+- `vectordb`: Defines a vectordb interface, contains several `Provider`s including `Memory`, `Chromem`, `Milvus`
+- `document` Defines a `Document` interface use for RAG, implemented `File`, `Http` document types. Provide a `Parser` interface which transform document content into specific string with `PDFParser` and `HTML to markdown` parsers implementations.
+
 3. `schema/`: Defines the Input/Output schema structures and interfaces
 4. `examples/`: Example projects showcasing Atomic Agents usage
 5. `tools/`: A collection of tools that can be used with Atomic Agents

components/document/doc.go

@@ -0,0 +1,2 @@
+// Package document contains Document structs and Parsers prepare for RAG
+package document

components/document/document.go

@@ -0,0 +1,35 @@
+package document
+
+import (
+	"bytes"
+	"errors"
+)
+
+var ErrReading = errors.New("document is reading")
+
+type ReadStatus = int32
+
+const (
+	Unread ReadStatus = iota
+	Reading
+	ReadCompleted
+)
+
+type ReadableDocument interface {
+	ReadAll() error
+	Read() (chan<- []byte, error)
+}
+
+type ClosableDocument interface {
+	Close() error
+}
+
+// Document is a document container with metadata
+type Document struct {
+	buffer *bytes.Buffer
+	Meta   map[string]string
+}
+
+func (d *Document) Reader() *bytes.Reader {
+	return bytes.NewReader(d.buffer.Bytes())
+}

components/document/file.go

@@ -0,0 +1,113 @@
+package document
+
+import (
+	"bytes"
+	"errors"
+	"io"
+	"os"
+	"strconv"
+
+	"go.uber.org/atomic"
+)
+
+type File struct {
+	status *atomic.Int32
+	fp     *os.File
+	Document
+}
+
+var (
+	_ ReadableDocument = (*File)(nil)
+	_ ClosableDocument = (*File)(nil)
+)
+
+func NewFile(fname string) (*File, error) {
+	fp, err := os.Open(fname)
+	if err != nil {
+		return nil, err
+	}
+	fileInfo, err := fp.Stat()
+	if err != nil {
+		return nil, err
+	}
+	if fileInfo.IsDir() {
+		return nil, errors.New("FileDocument could not be a directory")
+	}
+	return &File{
+		status: atomic.NewInt32(Unread),
+		fp:     fp,
+		Document: Document{
+			buffer: new(bytes.Buffer),
+			Meta: map[string]string{
+				"filename": fileInfo.Name(),
+				"modtime":  strconv.FormatInt(fileInfo.ModTime().Unix(), 10),
+			},
+		},
+	}, nil
+}
+
+func (d *File) ReadStatus() ReadStatus {
+	return d.status.Load()
+}
+
+func (d *File) ReadAll() error {
+	if d.ReadStatus() == Reading {
+		return ErrReading
+	} else if d.ReadStatus() == ReadCompleted {
+		return nil
+	}
+	if _, err := io.Copy(d.buffer, d.fp); err != nil {
+		d.status.Store(Unread)
+		return err
+	}
+	d.status.Store(ReadCompleted)
+	return nil
+}
+
+func (d *File) Read() (chan<- []byte, error) {
+	ch := make(chan<- []byte)
+	if d.ReadStatus() == Reading {
+		return nil, ErrReading
+	} else if d.ReadStatus() == ReadCompleted {
+		go func() {
+			defer close(ch)
+			d.status.Store(Reading)
+			reader := bytes.NewReader(d.buffer.Bytes())
+			tmp := make([]byte, 1024)
+			for {
+				n, err := reader.Read(tmp)
+				if err != nil {
+					d.status.Store(ReadCompleted)
+					return
+				}
+				bs := make([]byte, n)
+				copy(bs, tmp[:n])
+				ch <- bs
+			}
+		}()
+		return ch, nil
+	}
+	go func() {
+		defer close(ch)
+		d.status.Store(Reading)
+		tmp := make([]byte, 1024)
+		for {
+			n, err := d.fp.Read(tmp)
+			if err != nil {
+				if errors.Is(err, io.EOF) {
+					d.status.Store(ReadCompleted)
+				} else {
+					d.buffer.Reset()
+					d.status.Store(Unread)
+				}
+				return
+			}
+			d.buffer.Write(tmp[:n])
+		}
+	}()
+	return ch, nil
+}
+
+func (d *File) Close() error {
+	return d.fp.Close()
+}

components/document/html2md_parser.go

@@ -0,0 +1,33 @@
+package document
+
+import (
+	"bytes"
+	"context"
+	"io"
+
+	htmltomarkdown "github.com/JohannesKaufmann/html-to-markdown/v2"
+	"github.com/JohannesKaufmann/html-to-markdown/v2/converter"
+)
+
+// HTML2MDParser is a parser which parse html content to markdown
+type HTML2MDParser struct {
+	opts []converter.ConvertOptionFunc
+}
+
+var _ Parser = (*PDFParser)(nil)
+
+func NewHTML2MDParser(opts ...converter.ConvertOptionFunc) *HTML2MDParser {
+	return &HTML2MDParser{
+		opts: opts,
+	}
+}
+
+// Parse try to parse a html content from a bytes.Reader into a markdown content then write to an io.Writer
+func (h *HTML2MDParser) Parse(ctx context.Context, reader *bytes.Reader, writer io.Writer) error {
+	bs, err := htmltomarkdown.ConvertReader(reader, h.opts...)
+	if err != nil {
+		return err
+	}
+	_, err = writer.Write(bs)
+	return err
+}

components/document/http.go

@@ -0,0 +1,154 @@
+package document
+
+import (
+	"bytes"
+	"errors"
+	"io"
+	"net/http"
+
+	"go.uber.org/atomic"
+)
+
+type Http struct {
+	status  *atomic.Int32
+	client  *http.Client
+	httpReq *http.Request
+	Document
+}
+
+var (
+	_ ReadableDocument = (*File)(nil)
+	_ ClosableDocument = (*File)(nil)
+)
+
+type HttpConfig struct {
+	client  *http.Client
+	link    string
+	method  string
+	payload io.Reader
+}
+
+type HttpOption func(*HttpConfig)
+
+func WithHttpMethod(method string) HttpOption {
+	return func(h *HttpConfig) {
+		h.method = method
+	}
+}
+
+func WithHttpURL(link string) HttpOption {
+	return func(h *HttpConfig) {
+		h.link = link
+	}
+}
+
+func WithPayload(payload io.Reader) HttpOption {
+	return func(h *HttpConfig) {
+		h.payload = payload
+	}
+}
+
+func WithHttpClient(client *http.Client) HttpOption {
+	return func(h *HttpConfig) {
+		h.client = client
+	}
+}
+
+func NewHttp(opts ...HttpOption) (*Http, error) {
+	var cfg HttpConfig
+	for _, opt := range opts {
+		opt(&cfg)
+	}
+	if cfg.client == nil {
+		cfg.client = http.DefaultClient
+	}
+	httpReq, err := http.NewRequest(cfg.method, cfg.link, cfg.payload)
+	if err != nil {
+		return nil, err
+	}
+	return &Http{
+		status:  atomic.NewInt32(Unread),
+		client:  cfg.client,
+		httpReq: httpReq,
+		Document: Document{
+			buffer: new(bytes.Buffer),
+			Meta: map[string]string{
+				"url":    cfg.link,
+				"method": cfg.method,
+			},
+		},
+	}, nil
+}
+
+func (h *Http) ReadStatus() ReadStatus {
+	return h.status.Load()
+}
+
+func (h *Http) ReadAll() error {
+	if h.ReadStatus() == Reading {
+		return ErrReading
+	} else if h.ReadStatus() == ReadCompleted {
+		return nil
+	}
+	httpResp, err := h.client.Do(h.httpReq)
+	if err != nil {
+		h.status.Store(Unread)
+		return err
+	}
+	defer httpResp.Body.Close()
+	if _, err = io.Copy(h.buffer, httpResp.Body); err != nil {
+		h.status.Store(Unread)
+	}
+	h.status.Store(ReadCompleted)
+	return nil
+}
+
+func (h *Http) Read() (chan<- []byte, error) {
+	ch := make(chan<- []byte)
+	if h.ReadStatus() == Reading {
+		return nil, ErrReading
+	} else if h.ReadStatus() == ReadCompleted {
+		go func() {
+			defer close(ch)
+			h.status.Store(Reading)
+			reader := bytes.NewReader(h.buffer.Bytes())
+			tmp := make([]byte, 1024)
+			for {
+				n, err := reader.Read(tmp)
+				if err != nil {
+					h.status.Store(ReadCompleted)
+					return
+				}
+				bs := make([]byte, n)
+				copy(bs, tmp[:n])
+				ch <- bs
+			}
+		}()
+		return ch, nil
+	}
+	go func() {
+		defer close(ch)
+		h.status.Store(Reading)
+		httpResp, err := h.client.Do(h.httpReq)
+		if err != nil {
+			h.status.Store(Unread)
+			return
+		}
+		defer httpResp.Body.Close()
+		tmp := make([]byte, 1024)
+		for {
+			n, err := httpResp.Body.Read(tmp)
+			if err != nil {
+				if errors.Is(err, io.EOF) {
+					h.status.Store(ReadCompleted)
+				} else {
+					h.buffer.Reset()
+					h.status.Store(Unread)
+				}
+				return
+			}
+			h.buffer.Write(tmp[:n])
+		}
+	}()
+	return ch, nil
+}

components/document/parser.go

@@ -0,0 +1,11 @@
+package document
+
+import (
+	"bytes"
+	"context"
+	"io"
+)
+
+type Parser interface {
+	Parse(context.Context, *bytes.Reader, io.Writer) error
+}