├── .flake8 ├── .gitignore ├── .markdownlint.json ├── README.md ├── cli ├── __init__.py ├── embed.py ├── frag_cli.py ├── init_command.py ├── store_init_command.py ├── test_settings_command.py └── utils.py ├── frag ├── __init__.py ├── completions │ ├── __init__.py │ ├── base_bot.py │ ├── interface_bot.py │ ├── prompter.py │ └── summarizer_bot.py ├── embeddings │ ├── __init__.py │ ├── get_embed_api.py │ ├── hf_embed_api.py │ ├── ingest │ │ └── ingest_url.py │ └── store.py ├── frag.py ├── settings │ ├── .fragrc_default.yaml │ ├── __init__.py │ ├── bot_model_settings.py │ ├── bots_settings.py │ ├── embed_api_settings.py │ ├── embed_settings.py │ └── settings.py ├── typedefs │ ├── __init__.py │ ├── bot_comms_types.py │ └── embed_types.py └── utils │ ├── __init__.py │ ├── console.py │ └── singleton.py ├── poetry.lock ├── pyproject.toml ├── scripts ├── build_docs.py └── serve_docs.py └── templates ├── interface.message.html ├── interface.system.html ├── summarizer.system.html └── summarizer.user.html /.flake8: -------------------------------------------------------------------------------- 1 | [flake8] 2 | max-line-length = 98 3 | extend-ignore = E203 4 | ignore-path = .git, .idea, .vscode, scripts 5 | 6 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *.so 5 | 6 | # C extensions 7 | *.c 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | 27 | # PyInstaller 28 | # Usually these files are written by a python script from a template 29 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 30 | *.manifest 31 | *.spec 32 | 33 | # Installer logs 34 | pip-log.txt 35 | pip-delete-this-directory.txt 36 | 37 | # Unit test / coverage reports 38 | htmlcov/ 39 | .tox/ 40 | .nox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | .pytest_cache/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | db.sqlite3 58 | db.sqlite3-journal 59 | 60 | # Flask stuff: 61 | instance/ 62 | .webassets-cache 63 | 64 | # Scrapy stuff: 65 | .scrapy 66 | 67 | # Sphinx documentation 68 | docs/_build/ 69 | 70 | # PyBuilder 71 | target/ 72 | 73 | # Jupyter Notebook 74 | .ipynb_checkpoints 75 | 76 | # IPython 77 | profile_default/ 78 | ipython_config.py 79 | 80 | # pyenv 81 | .python-version 82 | 83 | # pipenv 84 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 85 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 86 | # having no cross-platform support, pipenv may update the lock file according to the current platform. 87 | # Under such scenarios, it is advised to exclude Pipfile.lock. 88 | #.Pipfile.lock 89 | 90 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 91 | __pypackages__/ 92 | 93 | # Celery stuff 94 | celerybeat-schedule 95 | celerybeat.pid 96 | 97 | # SageMath parsed files 98 | *.sage.py 99 | 100 | # Environments 101 | .env 102 | .env.local 103 | .venv 104 | venv/ 105 | ENV/ 106 | env.bak/ 107 | venv.bak/ 108 | 109 | # Spyder project settings 110 | .spyderproject 111 | .spyproject 112 | 113 | # Rope project settings 114 | .ropeproject 115 | 116 | # mkdocs documentation 117 | /site 118 | 119 | # mypy 120 | .mypy_cache/ 121 | 122 | # Pyre type checker 123 | .pyre/ 124 | 125 | # macOS files 126 | .DS_Store 127 | 128 | # .old directory 129 | .old/ 130 | 131 | .vscode 132 | db/ 133 | custom_db_path/ 134 | .frag/ 135 | -------------------------------------------------------------------------------- /.markdownlint.json: -------------------------------------------------------------------------------- 1 | { 2 | "MD013": false, 3 | "MD041": false 4 | } 5 | 6 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # FRAG: Focused Retrieval-Augmented Generation 2 | 3 | FRAG is a novel approach to retrieval-augmented generation that addresses key limitations of existing RAG (Retrieval Augmented Generation) systems. By decoupling the retrieval and generation phases, FRAG enables more flexible, efficient, and context-aware generation from large unstructured knowledge sources. 4 | 5 | It is part of a process of rewriting and modularisation for [raft](https://github.com/lumpenspace/raft). 6 | 7 | ## Motivation 8 | 9 | While RAG systems have shown promise in knowledge-intensive tasks, they suffer from several limitations: 10 | 11 | 1. **Overconfidence**: RAG models often overestimate the relevance of retrieved results, taking them for granted without critically assessing their applicability to the current context. 12 | 2. **Stylistic inconsistency**: Retrieved passages can have varying styles that differ from the main generation model, leading to inconsistent outputs. Rephrasing sources before storage can mitigate this but may result in information loss. 13 | 3. **Extraneous information**: Retrieved results, while relevant, may be overly long and contain extraneous details, polluting the context and hindering focused generation. 14 | 4. **Inefficiency**: Techniques like scratchpad generation are token-intensive and introduce response delays. 15 | 16 | FRAG addresses these issues by introducing a focused retrieval phase that identifies concise, relevant knowledge fragments, followed by a generation phase that attends to these fragments to produce coherent, context-aware outputs. 17 | 18 | ## Approach 19 | 20 | ### Summary 21 | 22 | Key components: 23 | 24 | - Embeddings Database 25 | - Interface Model 26 | - Multiple Archivists 27 | 28 | Flow: 29 | 30 | 1. Fetch closest N fragments meeting a threshold from embeddings db. 31 | 2. Perform focused retrieval via Archivists 32 | 3. Archivists summarize helpful fragments 33 | 4. Feed summaries to Interface Model for context-aware generation 34 | 5. Generate final response 35 | 36 | ```mermaid 37 | graph TD 38 | A[User Input] --> |Last question| E[context] 39 | B[Vector DB] --> |N closest fragments| C[Archivist 1] 40 | B --> |N closest fragments| D[Archivist 2] 41 | B --> |N closest fragments| F[Archivist N] 42 | C --> |Relevant summary| G[Merged summary] 43 | D --> |Relevant summary| G 44 | F --> |Relevant summary| G 45 | A --> E[Context] 46 | E --> |Previous interaction| B 47 | E --> H[Interface Model] 48 | G --> H 49 | ``` 50 | 51 | ### Architecture 52 | 53 | FRAG is composed by three main elements: 54 | 55 | - An embeddings databases containing the RAG corpus, as-is. 56 | - An Interface Model, which will interact directly with the user. This should be the most powerful and capable model 57 | - One or more Archivist 58 | 59 | For each interaction, the system will fetch the closest N fragments (over a proximity threshold P) to the last model and user outputs. 60 | 61 | At this point, there will be two generation phases 62 | 63 | 1. **Focused Retrieval**: 64 | - N instances of the Archivist are instantiated. Each is given the previous interaction, current question, and is briefed with determining whether the retrieved document can help answer the question. 65 | - In case of negative response, that Archivist won't relay an answer. 66 | - Otherwise, it will provide a summary of the document targeted specifically for the Persona embodied by the Interface Model to answer the last question. 67 | - A pointer to the source document is also added to the response. 68 | 69 | 2. **Context-Aware Generation**: 70 | - Retrieved and summarized fragments are fed as additional context to the Interface Model 71 | - The model attends to these focused knowledge snippets to inform its generation process 72 | - Fragments serve as an external memory that the model can selectively draw upon based on the current context 73 | 74 | By decoupling retrieval, selection and generation, FRAG allows for the use of task-specific retrieval and summarization models that can be optimized independently. The focused nature of the retrieved snippets helps to mitigate issues of overconfidence and extraneous information, while the notes' conciseness allows us to leave them within the context of the Interface Model for several more turns. 75 | 76 | ## Funding 77 | 78 | FRAG's development has been funded through a research grant from [Nous Research](https://github.com/nousresearch) 79 | 80 | ## License 81 | 82 | MIT 83 | -------------------------------------------------------------------------------- /cli/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | The Frag CLI interface. 3 | """ 4 | 5 | from .frag_cli import main 6 | 7 | if __name__ == "__main__": 8 | main() 9 | -------------------------------------------------------------------------------- /cli/embed.py: -------------------------------------------------------------------------------- 1 | from frag.utils import console 2 | 3 | 4 | def main() -> None: 5 | console.print("embed") 6 | -------------------------------------------------------------------------------- /cli/frag_cli.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | The main CLI entrypoint for frag. 4 | """ 5 | 6 | import click 7 | from click.core import Group 8 | 9 | from .init_command import main as init 10 | from .test_settings_command import main as test_settings 11 | from .store_init_command import main as store_init 12 | 13 | 14 | @click.group() 15 | def frag() -> None: 16 | """ 17 | Main CLI group for frag. 18 | """ 19 | 20 | 21 | frag.add_command(init) 22 | frag.add_command(test_settings) 23 | frag.add_command(store_init) 24 | 25 | 26 | main: Group = frag 27 | -------------------------------------------------------------------------------- /cli/init_command.py: -------------------------------------------------------------------------------- 1 | import os 2 | import click 3 | 4 | import jinja2 5 | import shutil 6 | from pathlib import Path 7 | 8 | from frag.settings import Settings, SettingsDict 9 | from frag.utils.console import console 10 | 11 | from .utils import create_or_override, section 12 | 13 | 14 | def init(path: str) -> SettingsDict | None: 15 | """ 16 | Initialize a new frag project 17 | 18 | Args: 19 | path (str): The path to the project's root 20 | """ 21 | sections = section() 22 | 23 | sections.section(title="Configuring project settings") 24 | 25 | dir_path: str | None = create_or_override(path=path, name=".frag", dir=True) 26 | 27 | if dir_path is None: 28 | return 29 | 30 | default_settings: str = Settings.defaults_path.read_text() 31 | 32 | console.log("[b]creating default config file[/b]") 33 | Path(dir_path, "config.yaml").write_text( 34 | jinja2.Template(default_settings).render( 35 | {"frag_dir": Path(dir_path).relative_to(path)} 36 | ) 37 | ) 38 | 39 | console.log(f"created config file: {Path(dir_path, 'config.yaml')}") 40 | 41 | db_path: str | None = create_or_override(path=dir_path, name="db", dir=True) 42 | if db_path is not None: 43 | console.log(f"created db path: {db_path}") 44 | 45 | console.log("[b]copying default templates[/b]") 46 | 47 | docstore_path: str | None = create_or_override( 48 | path=dir_path, name="docstore", dir=True 49 | ) 50 | if docstore_path is not None: 51 | console.log(f"created docstore path: {docstore_path}") 52 | 53 | console.log("[b]copying default templates[/b]") 54 | templates_src_path: Path = Path(__file__).parent.parent / "frag" / "templates" 55 | templates_dest_path: Path = Path(dir_path) / "templates" 56 | 57 | for src_file in templates_src_path.glob("*"): 58 | dest_file: Path = templates_dest_path / src_file.name 59 | if dest_file.exists(): 60 | dest_file.unlink() 61 | shutil.copy(src_file, dest_file) 62 | 63 | 64 | @click.command("init") 65 | @click.argument( 66 | "path", 67 | default=os.getcwd(), 68 | type=click.Path(exists=True, file_okay=False, dir_okay=True), 69 | ) 70 | def main(path: str) -> None: 71 | init(path=path) 72 | -------------------------------------------------------------------------------- /cli/store_init_command.py: -------------------------------------------------------------------------------- 1 | import click 2 | from frag.utils import console 3 | from pathlib import Path 4 | from typing import List, Tuple 5 | from frag.settings import Settings 6 | from frag.embeddings.store import EmbeddingStore 7 | 8 | 9 | def init_store(items: Tuple[str]) -> None: 10 | # Load settings 11 | settings: Settings = Settings.from_path() 12 | console.log(f"Settings: {settings.model_dump()}") 13 | # Initialize the embedding store 14 | store: EmbeddingStore = EmbeddingStore.instance or EmbeddingStore.create( 15 | settings.embeds 16 | ) 17 | 18 | urls: List[str] = [] 19 | paths: List[str] = [] 20 | 21 | # Determine if each item is a URL or a file path 22 | for item in items: 23 | if item.startswith("http://") or item.startswith("https://"): 24 | urls.append(item) 25 | elif Path(item).exists(): 26 | paths.append(item) 27 | else: 28 | console.log(f"Invalid item detected, neither URL nor path: {item}") 29 | 30 | console.log(f"URLs: {urls}") 31 | console.log(f"Paths: {paths}") 32 | console.log(f"Embedding store: {store}") 33 | 34 | if len(urls) > 0: 35 | from frag.embeddings.ingest.ingest_url import URLIngestor 36 | 37 | ingestor = URLIngestor(store=store) 38 | ingestor.ingest(urls) 39 | 40 | # Further processing can be added here to handle URLs and paths with the store 41 | 42 | 43 | @click.command("init:store") 44 | @click.argument("items", nargs=-1, type=str) 45 | def main(items: Tuple[str]) -> None: 46 | init_store(items=items) 47 | -------------------------------------------------------------------------------- /cli/test_settings_command.py: -------------------------------------------------------------------------------- 1 | import os 2 | import click 3 | import json 4 | 5 | from pathlib import Path 6 | from rich.prompt import Prompt 7 | from rich.panel import Panel 8 | from rich.pretty import Pretty 9 | 10 | from .utils import C 11 | 12 | from frag.settings import Settings 13 | from frag.utils.console import console, error_console 14 | 15 | 16 | def test_settings( 17 | path: str, json_string: str | None = None, return_settings: bool = False 18 | ) -> Settings | None: 19 | """ 20 | Test the settings class 21 | 22 | Args: 23 | path (str): The path to the config file to test 24 | json_string (str, optional): The json string to test 25 | """ 26 | settings: Settings | None = None 27 | if json_string is not None: 28 | settings = Settings.from_dict(json.loads(json_string)) 29 | else: 30 | if Path(path).exists(): 31 | 32 | settings = Settings.from_path(path) 33 | else: 34 | console.print( 35 | f"🤷 [{C.WARNING.value}]Settings path does not exist[/]:[dark_orange] {path} [/]\n", 36 | emoji=True, 37 | ) 38 | 39 | overwrite: str = Prompt.ask( 40 | "👷 Create a new config?", choices=["y", "n"], default="y" 41 | ) 42 | 43 | if overwrite != "y": 44 | error_console.log( 45 | f"[bold {C.ERROR.value}]Test cancelled. Create a new config with\ 46 | [code]frag init[/code][/]" 47 | ) 48 | return 49 | else: 50 | console.log("Creating a new config...") 51 | from .init_command import init 52 | 53 | init(path=path) 54 | 55 | if settings: 56 | console.log(Panel(Pretty(settings.model_dump()), title="Settings")) 57 | if return_settings: 58 | return settings 59 | 60 | 61 | @click.command("test:settings") 62 | @click.argument( 63 | "path", 64 | default=Path(os.getcwd(), ".frag"), 65 | type=click.Path(file_okay=False, dir_okay=True), 66 | ) 67 | @click.option("--json_string", "-j", default=None, type=click.STRING) 68 | def main(path: str, json_string: str | None = None) -> None: 69 | test_settings(path=path, json_string=json_string) 70 | -------------------------------------------------------------------------------- /cli/utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Utilities for the CLI 3 | """ 4 | 5 | import os 6 | from enum import Enum 7 | from typing import Self 8 | from rich.prompt import Prompt 9 | from rich.panel import Panel 10 | from frag.utils.console import console 11 | 12 | 13 | class C(Enum): 14 | """ 15 | Colors for the CLI 16 | """ 17 | 18 | PINK = "pink" 19 | TURQUOISE = "turquoise" 20 | GREEN = "green" 21 | RED = "red" 22 | YELLOW = "yellow" 23 | BLUE = "blue" 24 | WHITE = "white" 25 | DULL = "dull" 26 | ERROR = "red" 27 | WARNING = "black on dark_orange" 28 | SUCCESS = "green" 29 | 30 | def __repr__(self) -> str: 31 | return str(self.value) 32 | 33 | 34 | def create_or_override(path: str, name: str, dir: bool = False) -> str | None: 35 | """ 36 | Create or override a file at the specified path. 37 | """ 38 | check_path: str = os.path.join(path, name) 39 | abs_path: str = os.path.abspath(check_path) 40 | if os.path.exists(check_path): 41 | console.log( 42 | f"[b]{'directory' if dir else 'file'} already present:[/b] {abs_path}" 43 | ) 44 | overwrite: str = Prompt.ask("Overwrite?", choices=["y", "n"], default="n") 45 | 46 | if overwrite != "y": 47 | console.log( 48 | f"[bold {C.ERROR.value}]Initialization cancelled. {name} already exists.[/]" 49 | ) 50 | return 51 | 52 | console.log(f"[b]Creating {check_path}[/b]") 53 | os.makedirs(path, exist_ok=True) 54 | if dir: 55 | os.makedirs(check_path, exist_ok=True) 56 | else: 57 | with open(check_path, "w", encoding="utf-8") as f: 58 | f.write("") 59 | console.print(f"[bold {C.GREEN.value}]Successfully created {name}[/]") 60 | return check_path 61 | 62 | 63 | class Sections: 64 | # singleton 65 | 66 | def __new__(cls) -> "Sections": 67 | if not hasattr(cls, "instance"): 68 | cls.instance: Self = super(Sections, cls).__new__(cls) 69 | return cls.instance 70 | 71 | index = 0 72 | subindex = 0 73 | 74 | @property 75 | def title_index(self) -> str: 76 | return f"{self.index + 1}." 77 | 78 | @property 79 | def subtitle_index(self) -> str: 80 | return f"{self.title_index}{self.subindex + 1}. " 81 | 82 | def section(self, title: str, subtitle: str | None = None) -> None: 83 | self.title(title) 84 | if subtitle: 85 | self.subtitle(subtitle) 86 | 87 | def subsection(self, title: str) -> None: 88 | self.subtitle(title) 89 | 90 | def title(self, title: str) -> None: 91 | console.print(Panel(f"[bold green]{self.title_index}[/][bold]{title}[/]")) 92 | self.subindex = 0 93 | self.index += 1 94 | 95 | def subtitle(self, title: str) -> None: 96 | console.print(f"[bold green]{self.subtitle_index}[/][bold]{title}[/]") 97 | self.subindex += 1 98 | 99 | def reset(self) -> None: 100 | self.index = 0 101 | self.subindex = 0 102 | 103 | 104 | section = Sections 105 | -------------------------------------------------------------------------------- /frag/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | .. include:: ../README.md 3 | """ 4 | 5 | from frag import typedefs 6 | from frag.frag import Frag 7 | from frag.embeddings.store import EmbeddingStore 8 | from frag.completions import Prompter 9 | from frag.utils import console, error_console, live 10 | 11 | from dotenv import load_dotenv 12 | 13 | load_dotenv() 14 | 15 | __all__: list[str] = [ 16 | "Frag", 17 | "typedefs", 18 | "EmbeddingStore", 19 | "Prompter", 20 | "console", 21 | "live", 22 | "error_console", 23 | ] 24 | -------------------------------------------------------------------------------- /frag/completions/__init__.py: -------------------------------------------------------------------------------- 1 | from .prompter import Prompter 2 | 3 | Prompter = Prompter 4 | -------------------------------------------------------------------------------- /frag/completions/base_bot.py: -------------------------------------------------------------------------------- 1 | """ 2 | Base API client, from which both the prompter and the summariser inherit. 3 | """ 4 | 5 | import os 6 | from typing import List, Dict, Any, Literal 7 | 8 | import jinja2 9 | from litellm.main import ModelResponse, completion 10 | 11 | from frag.typedefs import MessageParam, Role, SystemMessage, UserMessage 12 | from frag.settings import BotModelSettings 13 | from frag.utils.console import error_console 14 | 15 | 16 | class BaseBot: 17 | """ 18 | Base API client class that handles interactions with the LLM model. 19 | """ 20 | 21 | settings: BotModelSettings 22 | client_type: Literal["interface", "summarizer"] | None = None 23 | system_template: jinja2.Template 24 | user_template: jinja2.Template 25 | messages: List[MessageParam] 26 | responder: str 27 | 28 | def __init__(self, settings: BotModelSettings, template_dir: str) -> None: 29 | """ 30 | Initializes the BaseApiClient with the given settings, template paths, and client type. 31 | 32 | :param settings: LLMSettings object containing configuration for the model. 33 | :param system_template_path: Path to the system message template file. 34 | :param user_template_path: Path to the user message template file. 35 | :param client_type: Type of the client, used to determine the template files. 36 | """ 37 | if self.client_type is None: 38 | raise ValueError("client_type must be provided") 39 | self.settings: BotModelSettings = settings 40 | self.load_templates(template_dir) 41 | 42 | def run( 43 | self, messages: List[MessageParam], **kwargs: Dict[str, Any] 44 | ) -> ModelResponse: 45 | """ 46 | Processes the given messages and performs completion using the LLM model. 47 | 48 | :param messages: List of ChatCompletionMessage objects to be processed. 49 | """ 50 | try: 51 | rendered_messages: List[MessageParam] = [ 52 | msg for msg in self._render(messages, **kwargs) 53 | ] 54 | 55 | return ModelResponse( 56 | completion( 57 | model=self.settings.api, 58 | messages=rendered_messages, 59 | **self.settings.completion_kwargs, 60 | ) 61 | ) 62 | except Exception as e: 63 | error_console.log(f"Error during completion: {e}") 64 | raise 65 | 66 | def _render( 67 | self, messages: List[MessageParam], **kwargs: Dict[str, Any] 68 | ) -> List[MessageParam]: 69 | """ 70 | Renders the messages based on the templates. This method should be implemented 71 | by subclasses. 72 | 73 | :param messages: List of ChatCompletionMessage objects to be rendered. 74 | :return: List of rendered ChatCompletionMessage objects. 75 | """ 76 | raise NotImplementedError 77 | 78 | def load_templates(self, template_dir: str) -> None: 79 | """ 80 | Loads the message templates from the specified paths. 81 | 82 | :param system_template_path: Path to the system message template file. 83 | :param user_template_path: Path to the user message template file. 84 | """ 85 | if self.client_type is None: 86 | raise ValueError("[dev] client_type must be set") 87 | 88 | template_dir = template_dir if template_dir else "templates" 89 | try: 90 | with open( 91 | os.path.join(template_dir, f"{self.client_type}.system.html"), 92 | "r", 93 | encoding="utf-8", 94 | ) as file: 95 | self.system_template = jinja2.Template(file.read()) 96 | with open( 97 | os.path.join(template_dir, f"{self.client_type}.user.html"), 98 | "r", 99 | encoding="utf-8", 100 | ) as file: 101 | self.user_template = jinja2.Template(file.read()) 102 | except FileNotFoundError as e: 103 | error_console.log("Template file not found: %s", e) 104 | raise e 105 | except Exception as e: 106 | error_console.log("Error loading templates: %s", e) 107 | raise e 108 | 109 | def _render_message( 110 | self, latest_messages: List[MessageParam], role: Role, **kwargs: Dict[str, Any] 111 | ) -> MessageParam: 112 | """ 113 | Renders a message based on the role and the latest messages. 114 | 115 | :param latest_messages: List of the most recent ChatCompletionMessage objects. 116 | :param role: Role of the message to be rendered (SYSTEM or USER). 117 | :return: Rendered ChatCompletionMessage object. 118 | """ 119 | if role == SystemMessage: 120 | return SystemMessage( 121 | content=self.system_template.render( 122 | latest_messages=latest_messages, **kwargs 123 | ), 124 | role="system", 125 | ) 126 | elif role == UserMessage: 127 | return UserMessage( 128 | content=self.user_template.render( 129 | latest_messages=latest_messages, **kwargs 130 | ), 131 | role="user", 132 | ) 133 | else: 134 | raise ValueError("MessageType must be SystemMessage or UserMessage") 135 | -------------------------------------------------------------------------------- /frag/completions/interface_bot.py: -------------------------------------------------------------------------------- 1 | """ 2 | Interface bot. This is the bot that will directly interact with the user. 3 | 4 | To personalise the template, extract (`frag template extract`) and modify them from: 5 | - .frag/templates/system.j2 6 | - .frag/templates/user.j2 7 | """ 8 | 9 | from typing import List 10 | 11 | import jinja2 12 | from litellm import completion, ModelResponse 13 | 14 | from frag.typedefs import MessageParam, Note 15 | from frag.settings import BotModelSettings 16 | from frag.utils.console import error_console 17 | 18 | from .base_bot import BaseBot 19 | 20 | 21 | class InterfaceBot(BaseBot): 22 | """ 23 | Interface bot. This is the bot that will directly interact with the user. 24 | 25 | It is called by `prompter.py`. 26 | """ 27 | 28 | client_type = "interface" 29 | settings: BotModelSettings 30 | system_template: jinja2.Template 31 | user_template: jinja2.Template 32 | 33 | def __init__(self, settings: BotModelSettings, template_dir: str) -> None: 34 | super().__init__(settings, template_dir=template_dir) 35 | 36 | def run(self, messages: List[MessageParam], notes: List[Note]) -> ModelResponse: 37 | try: 38 | rendered_messages: List[MessageParam] = self._render(messages, notes=notes) 39 | 40 | result = ModelResponse( 41 | completion( 42 | messages=rendered_messages, 43 | model=self.settings.api, 44 | stream=False, 45 | **self.settings.model_dump(exclude=set(["llm"])), 46 | ) 47 | ) 48 | return result 49 | except jinja2.TemplateError as te: 50 | error_console.log("Template rendering error: %s", te) 51 | raise 52 | except Exception as e: 53 | error_console.log("General error during completion: %s", e) 54 | raise 55 | 56 | def _render( 57 | self, messages: List[MessageParam], **kwargs: MessageParam 58 | ) -> List[MessageParam]: 59 | try: 60 | last_message = messages[-1] 61 | return [ 62 | self._render_message(messages=messages[:-1], role="system", **kwargs), 63 | *messages[:-1], 64 | self._render_message( 65 | messages=last_message.get("content"), role="user", **kwargs 66 | ), 67 | ] 68 | except IndexError as ie: 69 | error_console.log( 70 | "Rendering error - possibly due to empty messages list: %s", ie 71 | ) 72 | raise 73 | except Exception as e: 74 | error_console.log("General rendering error: %s", e) 75 | raise 76 | -------------------------------------------------------------------------------- /frag/completions/prompter.py: -------------------------------------------------------------------------------- 1 | from typing import List, Any 2 | from litellm import Choices, ModelResponse 3 | from frag.settings import BotsSettings 4 | from frag.typedefs import MessageParam 5 | from .summarizer_bot import SummarizerBot 6 | from .interface_bot import InterfaceBot 7 | from frag.utils.console import error_console 8 | 9 | 10 | class Prompter: 11 | """ 12 | Main handler of prompts and responses. 13 | 14 | """ 15 | 16 | def __init__( 17 | self, settings: BotsSettings, summarizer: SummarizerBot, interface: InterfaceBot 18 | ) -> None: 19 | self.settings: BotsSettings = settings 20 | self.summarizer: SummarizerBot = summarizer 21 | self.interface: InterfaceBot = interface 22 | 23 | def respond(self, messages: List[MessageParam], **kwargs: Any) -> str: 24 | """ 25 | Respond to a message history. 26 | """ 27 | try: 28 | responses: List[Choices] = [ 29 | Choices(c) for c in self.interface.run(messages, **kwargs).choices 30 | ] 31 | if responses and responses[0]: 32 | return str(responses[0].get("content", "")) 33 | return "" 34 | except Exception as e: 35 | error_console.log("Error in responding: %s", e) 36 | raise 37 | 38 | def summarise(self, messages: List[MessageParam], **kwargs: Any) -> ModelResponse: 39 | try: 40 | return self.summarizer.run(messages, **kwargs) 41 | except Exception as e: 42 | error_console.log("Error in summarising: %s", e) 43 | raise 44 | -------------------------------------------------------------------------------- /frag/completions/summarizer_bot.py: -------------------------------------------------------------------------------- 1 | """ 2 | The SummarizerBot class is responsible for summarizing chat messages. 3 | It uses internal methods to render system and user messages based on 4 | the latest messages. 5 | """ 6 | 7 | from typing import List, Dict, Any 8 | from llama_index.core.schema import Document 9 | from frag.settings.bot_model_settings import BotModelSettings 10 | from frag.typedefs import MessageParam, DocMeta 11 | from .base_bot import BaseBot 12 | 13 | 14 | class SummarizerBot(BaseBot): 15 | """ 16 | The Summarizer class is responsible for summarizing chat messages. 17 | It uses internal methods to render system and user messages based on 18 | the latest messages. 19 | """ 20 | 21 | client_type = "summarizer" 22 | 23 | def __init__(self, settings: BotModelSettings, template_dir: str) -> None: 24 | super().__init__(settings, template_dir=template_dir) 25 | 26 | def _render( 27 | self, 28 | *_, 29 | messages: List[MessageParam], 30 | document: Document, 31 | doc_meta: DocMeta, 32 | **kwargs: Dict[str, Any] 33 | ) -> List[MessageParam]: 34 | return [ 35 | self._render_message(messages, role="system", **kwargs), 36 | self._render_message(messages, role="user", **kwargs), 37 | ] 38 | -------------------------------------------------------------------------------- /frag/embeddings/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lumpenspace/FRAG/bad1d560047b45d51a6014a12cefede2e641e443/frag/embeddings/__init__.py -------------------------------------------------------------------------------- /frag/embeddings/get_embed_api.py: -------------------------------------------------------------------------------- 1 | from llama_index.core.embeddings import BaseEmbedding 2 | from frag.typedefs.embed_types import ApiSource 3 | 4 | 5 | def get_embed_api( 6 | api_source: ApiSource, api_model: str | None, api_key: str | None = None 7 | ) -> BaseEmbedding: 8 | """ 9 | Retrieves an embedding API instance based on the input. 10 | 11 | Args: 12 | embed_api (EmbedAPI|str): The embedding API instance or a string identifier for the API. 13 | For OpenAI APIs, prepend 'oai:' to the API name. 14 | 15 | Returns: 16 | EmbedAPI: An instance of the requested embedding API. 17 | 18 | Raises: 19 | ValueError: If no embedding API is provided. 20 | """ 21 | 22 | if not api_model: 23 | raise ValueError("Embedding model and chunking settings must be provided") 24 | if api_source == "OpenAI": 25 | from llama_index.embeddings.openai import OpenAIEmbedding 26 | 27 | return OpenAIEmbedding(model=api_model, api_key=api_key) 28 | elif api_source == "HuggingFace": 29 | from llama_index.embeddings.huggingface import HuggingFaceEmbedding 30 | 31 | try: 32 | embed_api = HuggingFaceEmbedding(model=api_model) 33 | except ValueError: 34 | raise ValueError(f"Invalid embedding API: {api_model} on {api_source}") 35 | 36 | raise ValueError("Invalid embedding API type: %s" % type(embed_api)) 37 | -------------------------------------------------------------------------------- /frag/embeddings/hf_embed_api.py: -------------------------------------------------------------------------------- 1 | """ 2 | This module allows for embedding with HuggingFace models. 3 | """ 4 | 5 | from typing import List 6 | from pydantic import Field 7 | from chromadb.api.types import EmbeddingFunction, Documents 8 | from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction 9 | from sentence_transformers import SentenceTransformer 10 | 11 | from llama_index.core.embeddings import BaseEmbedding 12 | 13 | 14 | class HFEmbedAPI(BaseEmbedding): 15 | """ 16 | A class to interact with HuggingFace's embedding models. 17 | 18 | Attributes: 19 | name: The name of the HuggingFace embedding model, for example "gpt2". 20 | max_tokens: The maximum number of tokens to embed. 21 | """ 22 | 23 | name: str = Field( 24 | "all-MiniLM-L6-v2", description="Name for HuggingFace embeddings model" 25 | ) 26 | max_tokens: int = Field(512, description="Maximum tokens to embed") 27 | 28 | _api: SentenceTransformer | None = None 29 | 30 | @property 31 | def model(self) -> SentenceTransformer: 32 | 33 | if self._api is None: 34 | try: 35 | self._api = SentenceTransformer(self.name) 36 | except ImportError: 37 | raise ImportError( 38 | "Unable to import SentenceTransformer. Please install it using ", 39 | "`pip install sentence-transformers`", 40 | ) 41 | except Exception as e: 42 | raise ValueError(f"Error embedding text with HF model: {e}") 43 | return self._api 44 | 45 | def encode(self, text: str) -> List[int]: 46 | return self.model.tokenizer.encode(text) 47 | 48 | def decode(self, tokens: List[int]) -> str: 49 | return self.model.tokenizer.decode(tokens) 50 | 51 | @property 52 | def embed_function(self) -> EmbeddingFunction[Documents]: 53 | try: 54 | return SentenceTransformerEmbeddingFunction(model_name=self.name) 55 | except Exception as e: 56 | raise ValueError(f"Error embedding text with HF model: {e}") 57 | -------------------------------------------------------------------------------- /frag/embeddings/ingest/ingest_url.py: -------------------------------------------------------------------------------- 1 | from typing import Callable, List, Dict, Any 2 | from frag.embeddings.store import EmbeddingStore, AddOns 3 | from llama_index.core.schema import Document 4 | from llama_index.readers.web import ( 5 | BeautifulSoupWebReader, 6 | ) 7 | from llama_index.core.ingestion import IngestionPipeline 8 | from pydantic import BaseModel, field_validator, Field, ConfigDict 9 | 10 | 11 | class URLIngestor(BaseModel): 12 | """ 13 | Ingest a URL into the embedding store. 14 | """ 15 | 16 | store: EmbeddingStore | None = Field(default=None) 17 | 18 | model_config: ConfigDict = ConfigDict( 19 | arbitrary_types_allowed=True, 20 | ) 21 | 22 | @property 23 | def pipeline_addons(self) -> AddOns: 24 | return { 25 | "extractors": [], 26 | "preprocessors": [], 27 | } 28 | 29 | def ingest(self, data: List[str] | str) -> None: 30 | """ 31 | Ingest the URL into the embedding store. 32 | """ 33 | if isinstance(data, str): 34 | data = [data] 35 | self.store.get_pipeline(self.pipeline_addons).run( 36 | documents=BeautifulSoupWebReader().load_data(urls=data), 37 | ) 38 | -------------------------------------------------------------------------------- /frag/embeddings/store.py: -------------------------------------------------------------------------------- 1 | from typing import List, TypedDict, Self 2 | 3 | from llama_index.core import StorageContext, VectorStoreIndex 4 | from llama_index.core.ingestion import IngestionCache, IngestionPipeline 5 | from llama_index.core.extractors import BaseExtractor 6 | from llama_index.core.node_parser import NodeParser 7 | from llama_index.core.schema import TransformComponent 8 | from llama_index.core.storage.docstore import SimpleDocumentStore 9 | from llama_index.core.retrievers import BaseRetriever 10 | from llama_index.core.node_parser import SentenceSplitter 11 | from llama_index.vector_stores.chroma import ChromaVectorStore 12 | from chromadb.api import ClientAPI 13 | from chromadb.api.models.Collection import Collection 14 | from chromadb import PersistentClient 15 | from trio import Path 16 | from frag.settings.embed_settings import EmbedSettings 17 | from frag.utils import SingletonMixin 18 | from frag.typedefs.embed_types import BaseEmbedding 19 | 20 | ArgType = TypedDict( 21 | "ArgType", 22 | {"settings": EmbedSettings}, 23 | ) 24 | 25 | AddOns = TypedDict( 26 | "AddOns", 27 | {"extractors": List[BaseExtractor], "preprocessors": List[TransformComponent]}, 28 | ) 29 | 30 | 31 | class EmbeddingStore(SingletonMixin[type(ArgType)]): 32 | embed_model: BaseEmbedding 33 | db: ClientAPI 34 | collection: Collection 35 | text_splitter: SentenceSplitter = SentenceSplitter() 36 | docstore: SimpleDocumentStore = SimpleDocumentStore() 37 | vector_store: ChromaVectorStore 38 | settings: EmbedSettings 39 | index: BaseRetriever 40 | collection_name: str 41 | text_splitter: SentenceSplitter 42 | docstore: SimpleDocumentStore 43 | vector_store: ChromaVectorStore 44 | 45 | def __init__( 46 | self, 47 | settings: EmbedSettings, 48 | collection_name: str | None = None, 49 | ) -> None: 50 | self.settings = settings 51 | self.collection_name = collection_name or self.settings.default_collection 52 | self.db = PersistentClient( 53 | path=str(settings.path / "db"), 54 | ) 55 | self.embed_model = settings.api 56 | self.change_collection(collection_name=collection_name) 57 | self.text_splitter = SentenceSplitter() 58 | self.docstore = SimpleDocumentStore() 59 | self.vector_store = ChromaVectorStore(chroma_collection=self.collection) 60 | 61 | @classmethod 62 | def create( 63 | cls, settings: EmbedSettings, collection_name: str | None = None 64 | ) -> Self: 65 | """ 66 | Create a new embedding store 67 | """ 68 | cls.reset() 69 | instance: Self = cls.__new__(cls, settings=settings) 70 | return instance 71 | 72 | def get_index(self) -> BaseRetriever: 73 | """ 74 | Get the index 75 | """ 76 | vector_store: ChromaVectorStore = ChromaVectorStore( 77 | chroma_collection=self.collection 78 | ) 79 | storage_context: StorageContext = StorageContext.from_defaults( 80 | vector_store=vector_store 81 | ) 82 | index: VectorStoreIndex = VectorStoreIndex.from_documents( 83 | documents=[], 84 | storage_context=storage_context, 85 | embed_model=self.embed_model, 86 | ) 87 | return index.as_retriever() 88 | 89 | def change_collection(self, collection_name: str | None = None) -> None: 90 | """ 91 | Change the collection name 92 | """ 93 | self.collection_name = collection_name or self.settings.default_collection 94 | self.collection = self.db.get_or_create_collection(name=self.collection_name) 95 | self.index = self.get_index() 96 | 97 | def get_pipeline( 98 | self, 99 | addons: AddOns, 100 | ) -> IngestionPipeline: 101 | """ 102 | Ingest a URL into the embedding store. 103 | """ 104 | return IngestionPipeline( 105 | transformations=[ 106 | *addons["preprocessors"], 107 | self.text_splitter, 108 | self.embed_model, 109 | *addons["extractors"], 110 | ], 111 | cache=IngestionCache( 112 | collection=f"{self.collection_name}-{self.embed_model.model_name}" 113 | ), 114 | vector_store=self.vector_store, 115 | docstore=self.docstore, 116 | ) 117 | -------------------------------------------------------------------------------- /frag/frag.py: -------------------------------------------------------------------------------- 1 | """ 2 | fRAG: Main class for the fRAG library 3 | """ 4 | 5 | from pathlib import Path 6 | from frag.embeddings.store import EmbeddingStore 7 | from frag.settings import Settings 8 | from frag.settings.settings import SettingsDict 9 | 10 | 11 | class Frag: 12 | """ 13 | Main class for the fRAG library 14 | """ 15 | 16 | def __init__(self, settings: Settings | SettingsDict | str) -> None: 17 | if isinstance(settings, str): 18 | if settings.startswith("~"): 19 | settings = settings.replace("~", str(Path.home())) 20 | settings = Settings.from_path(settings) 21 | if isinstance(settings, SettingsDict): 22 | settings = Settings.from_dict(settings) 23 | self.settings: Settings = settings 24 | self.embedding_store: EmbeddingStore = EmbeddingStore.create( 25 | settings=self.settings.embeds, 26 | ) 27 | -------------------------------------------------------------------------------- /frag/settings/.fragrc_default.yaml: -------------------------------------------------------------------------------- 1 | embeds: 2 | default_collection: default # name of the default collection 3 | api_model: text-embedding-3-large # Name of the embedding model. 4 | # for HuggingFace, 5 | # see https://www.sbert.net/docs/pretrained_models.html 6 | # for OpenAI, see https://platform.openai.com/docs/models 7 | api_source: OpenAI # source of the embedding model. can be OpenAI or HuggingFace 8 | # also available: 9 | # max_tokens(int), to set the maximum number of tokens to embed 10 | bots: 11 | api: gpt-3.5-turbo # see: https://litellm.vercel.app/docs/providers 12 | # we use the lite-llm default settings unless the user specifies otherwise, 13 | interface: { api: gpt-4-turbo } # these settings override top-level settings for the interface bot 14 | # you can also specify other bot-specific settings, e.g. 15 | # summarizer: { max_tokens: 200 } 16 | # extractor: { api: gpt-3.5-turbo } 17 | 18 | -------------------------------------------------------------------------------- /frag/settings/__init__.py: -------------------------------------------------------------------------------- 1 | from .bots_settings import BotsSettings, BotModelSettings 2 | from .settings import Settings, SettingsDict 3 | from .embed_settings import EmbedSettings, EmbedSettingsDict 4 | 5 | __all__: list[str] = [ 6 | "Settings", 7 | "SettingsDict", 8 | "EmbedSettings", 9 | "EmbedSettingsDict", 10 | "BotsSettings", 11 | "BotModelSettings", 12 | ] 13 | -------------------------------------------------------------------------------- /frag/settings/bot_model_settings.py: -------------------------------------------------------------------------------- 1 | """ 2 | LLM settings, used for both the interface and summarizer bots. 3 | """ 4 | 5 | from typing import Any, Dict 6 | from litellm import get_model_info, get_supported_openai_params 7 | from pydantic import model_validator 8 | from pydantic_settings import BaseSettings 9 | 10 | 11 | class BotModelSettings(BaseSettings): 12 | """ 13 | LLM settings, used for both the interface and summarizer bots. 14 | """ 15 | 16 | api: str = "gpt-3.5-turbo" 17 | bot: str 18 | 19 | def dump(self) -> dict[str, Any]: 20 | """ 21 | return model settings 22 | """ 23 | return self.model_dump(exclude_unset=True, exclude_none=True) 24 | 25 | def __getitem__(self, item: str) -> Any: 26 | return self.model_dump(exclude_unset=True, exclude_none=True)[item] 27 | 28 | @property 29 | def api_name(self) -> str: 30 | """ 31 | Returns the API name for the model. 32 | """ 33 | return self.model_dump(exclude_unset=True, exclude_none=True)["api"] 34 | 35 | @model_validator(mode="before") 36 | @classmethod 37 | def validate_model(cls, values: dict[str, Any]) -> dict[str, Any]: 38 | """ 39 | Makes sure all params are supported by the model. 40 | """ 41 | other_values: Dict[str, Any] = { 42 | k: v for k, v in values.items() if k not in ["api", "bot"] 43 | } 44 | # check if values are supported 45 | model_string: str = values.get("api", "gpt-3.5-turbo") 46 | 47 | try: 48 | info: Dict[str, Any] = get_model_info(model_string) 49 | except Exception: 50 | raise ValueError( 51 | f"Error getting model info for {model_string}.\n\ 52 | List of supported models:\n\ 53 | https://docs.litellm.ai/docs/providers\n" 54 | ) 55 | supported_params: list[str] | None = get_supported_openai_params( 56 | model_string, info.get("litellm_provider", None) 57 | ) 58 | if supported_params is None: 59 | raise ValueError("Model {model_string} is not supported by litellm.") 60 | for k, _ in other_values.items(): 61 | if k not in supported_params: 62 | raise ValueError( 63 | f"Unsupported parameter {k} for model {values.get('model', 'gpt-3.5-turbo')}" 64 | ) 65 | return { 66 | **other_values, 67 | "api": model_string, 68 | "bot": values.get("bot"), 69 | } 70 | -------------------------------------------------------------------------------- /frag/settings/bots_settings.py: -------------------------------------------------------------------------------- 1 | """ 2 | LLM settings, used for both the interface and summarizer bots. 3 | 4 | In .frag/config.yaml, the settings are under `bots` - default ones direct descendants, and bot 5 | specifics under the bot name. For example, here: 6 | 7 | ```yaml 8 | bots: 9 | api: gpt-3.5-turbo 10 | max_tokens: 1000 11 | top_p: 1.0 12 | interface: 13 | api: gpt-4-turbo 14 | summarizer: 15 | temperature: 0.5 16 | extractor: 17 | max_tokens: 200 18 | ``` 19 | 20 | ... the `model` key is the default model for all bots, and the `interface` and `summarizer` keys 21 | are bot specific settings. If a bot specific setting is not present, the default one is used. 22 | """ 23 | 24 | from typing import Dict, Any, Self 25 | from frag.utils.console import console 26 | from .bot_model_settings import BotModelSettings 27 | from pydantic_settings import BaseSettings 28 | 29 | 30 | class BotsSettings(BaseSettings): 31 | """ 32 | LLM settings, used for both the interface and summarizer bots. 33 | """ 34 | 35 | interface_bot: BotModelSettings 36 | summarizer_bot: BotModelSettings 37 | extractor_bot: BotModelSettings 38 | 39 | @classmethod 40 | def from_dict( 41 | cls, 42 | bots_dict: Dict[str, Any], 43 | ) -> Self: 44 | """ 45 | Create an LLMSettings object from a dictionary. 46 | """ 47 | interface_bot_settings: Dict[str, Any] = bots_dict.pop("interface", {}) 48 | summarizer_bot_settings: Dict[str, Any] = bots_dict.pop("summarizer", {}) 49 | extractor_bot_settings: Dict[str, Any] = bots_dict.pop("extractor", {}) 50 | default_settings: Dict[str, Any] = bots_dict 51 | 52 | console.log( 53 | f"[b]BotsSettings[/b]\ 54 | \n\tdefault: {default_settings}\ 55 | \n\tinterface: {interface_bot_settings}\ 56 | \n\tsummarizer: {summarizer_bot_settings}\ 57 | \n\textractor: {extractor_bot_settings}" 58 | ) 59 | 60 | return cls( 61 | interface_bot=BotModelSettings( 62 | # merge the default settings with the bot-specific settings 63 | # and add the bot name to the settings 64 | **{**default_settings, **interface_bot_settings}, 65 | bot="interface", 66 | ), 67 | summarizer_bot=BotModelSettings( 68 | **{**default_settings, **summarizer_bot_settings}, bot="summarizer" 69 | ), 70 | extractor_bot=BotModelSettings( 71 | **{**default_settings, **extractor_bot_settings}, bot="extractor" 72 | ), 73 | ) 74 | -------------------------------------------------------------------------------- /frag/settings/embed_api_settings.py: -------------------------------------------------------------------------------- 1 | from typing import Self, Dict, Any 2 | from typing_extensions import TypedDict 3 | from pydantic_settings import BaseSettings 4 | 5 | from frag.embeddings.get_embed_api import get_embed_api 6 | from llama_index.core.embeddings import BaseEmbedding 7 | from frag.typedefs.embed_types import ApiSource 8 | from frag.utils import console 9 | 10 | EmbedApiSettingsDict = TypedDict( 11 | "EmbedApiSettingsDict", 12 | {"api_name": str, "api_source": ApiSource, "chunk_overlap": int | None}, 13 | ) 14 | 15 | 16 | class EmbedAPISettings(BaseSettings): 17 | api_source: ApiSource = "OpenAI" 18 | api_model: str = "text-embedding-3-large" 19 | 20 | chunk_overlap: int = 0 21 | 22 | @property 23 | def api(self) -> BaseEmbedding: 24 | if not hasattr(self, "_api"): 25 | self._api: BaseEmbedding = get_embed_api( 26 | api_model=self.api_model, api_source=self.api_source 27 | ) 28 | return self._api 29 | 30 | @classmethod 31 | def from_dict( 32 | cls, 33 | embeds_dict: Dict[str, Any], 34 | ) -> Self: 35 | console.log(f"[b]EmbedApiSettings:[/] {embeds_dict}") 36 | api_model: str = embeds_dict.get("api_model", "text-embedding-3-large") 37 | api_source: ApiSource = embeds_dict.get("api_source", "OpenAI") 38 | 39 | instance = cls( 40 | api_model=api_model, 41 | api_source=api_source, 42 | chunk_overlap=embeds_dict.get("chunk_overlap", 0), 43 | ) 44 | 45 | try: 46 | instance.api 47 | except Exception as e: 48 | console.log("Error getting embed api", e) 49 | 50 | return instance 51 | -------------------------------------------------------------------------------- /frag/settings/embed_settings.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | from typing import Self, Dict, Any 3 | from typing_extensions import TypedDict 4 | from pydantic_settings import BaseSettings 5 | from pydantic import field_validator, ValidationInfo 6 | import logging 7 | 8 | from frag.embeddings.get_embed_api import get_embed_api 9 | from llama_index.core.embeddings import BaseEmbedding 10 | from frag.typedefs.embed_types import ApiSource 11 | from frag.utils.console import console, error_console 12 | 13 | EmbedSettingsDict = TypedDict( 14 | "EmbedSettingsDict", 15 | { 16 | "api_name": str, 17 | "api_source": ApiSource, 18 | "chunk_overlap": int, 19 | "path": Path, 20 | "default_collection": str, 21 | }, 22 | ) 23 | 24 | 25 | class EmbedSettings(BaseSettings): 26 | api_source: ApiSource = "OpenAI" 27 | api_model: str = "text-embedding-3-large" 28 | chunk_overlap: int = 0 29 | path: Path = Path("./db") 30 | default_collection: str = "default" 31 | 32 | @field_validator("default_collection") 33 | @classmethod 34 | def validate_default_collection(cls, v: str) -> str: 35 | """ 36 | Validate the collection name. 37 | """ 38 | if not v: 39 | v = "default_collection" 40 | logging.warning( 41 | "Collection name not provided, using default collection name: %s", 42 | v, 43 | ) 44 | return v 45 | 46 | @property 47 | def api(self) -> BaseEmbedding: 48 | if not hasattr(self, "_api"): 49 | self._api: BaseEmbedding = get_embed_api( 50 | api_model=self.api_model, api_source=self.api_source 51 | ) 52 | return self._api 53 | 54 | @classmethod 55 | def from_dict( 56 | cls, 57 | embeds_dict: Dict[str, Any], 58 | ) -> Self: 59 | console.log(f"[b]EmbedApiSettings:[/] {embeds_dict}") 60 | api_model: str = embeds_dict.get("api_model", "text-embedding-3-large") 61 | api_source: ApiSource = embeds_dict.get("api_source", "OpenAI") 62 | instance: Self = cls( 63 | api_model=api_model, 64 | api_source=api_source, 65 | chunk_overlap=embeds_dict.get("chunk_overlap", 0), 66 | default_collection=embeds_dict.get("default_collection", "default"), 67 | path=Path(embeds_dict.get("path", "./db")), 68 | ) 69 | try: 70 | instance.api 71 | except Exception as e: 72 | error_console.log("Error getting embed api", e) 73 | 74 | return instance 75 | -------------------------------------------------------------------------------- /frag/settings/settings.py: -------------------------------------------------------------------------------- 1 | """ 2 | Full settings for the frag system. 3 | 4 | They are by default read from a .fragrc file, which you can generate by running: 5 | 6 | ```sh 7 | frag init 8 | ``` 9 | """ 10 | 11 | import os 12 | from typing import Dict, Any, Self 13 | from typing_extensions import TypedDict 14 | from pathlib import Path 15 | import yaml 16 | import json 17 | from pydantic_settings import BaseSettings 18 | 19 | from frag.utils.singleton import SingletonMixin 20 | 21 | from .bots_settings import BotsSettings 22 | from .embed_settings import EmbedSettings, EmbedSettingsDict 23 | from frag.utils.console import console, error_console, live 24 | 25 | 26 | EmbedSettingsDict = EmbedSettingsDict 27 | 28 | SettingsDict = TypedDict( 29 | "SettingsDict", 30 | { 31 | "embeds": None | EmbedSettingsDict, 32 | "bots": None | Dict[str, Any], 33 | "path": None | Path, 34 | }, 35 | ) 36 | 37 | 38 | class Settings(BaseSettings, SingletonMixin[type(SettingsDict)]): 39 | """ 40 | Read settings from a .frag file or load them programmatically, and validate them. 41 | 42 | those settings include: 43 | """ 44 | 45 | embeds: EmbedSettings 46 | bots: BotsSettings 47 | path: Path 48 | __pickled__: bool = False 49 | 50 | _frag_dir: Path | None = None 51 | 52 | @classmethod 53 | def create(cls, embeds: EmbedSettings, bots: BotsSettings, path: Path) -> Self: 54 | return cls(embeds=embeds, bots=bots, path=path) 55 | 56 | @classmethod 57 | def set_dir(cls, frag_dir: str) -> Path: 58 | """ 59 | Set the directory in which the fragrc file is located. 60 | """ 61 | cls._frag_dir = Path(Path.cwd(), frag_dir) 62 | return cls._frag_dir 63 | 64 | @classmethod 65 | @property 66 | def defaults_path(cls) -> Path: 67 | dir_path: str = os.path.dirname(os.path.realpath(__file__)) 68 | return Path(dir_path) / ".fragrc_default.yaml" 69 | 70 | @classmethod 71 | @property 72 | def frag_dir(cls) -> Path: 73 | """ 74 | Get the directory in which the fragrc file is located. 75 | """ 76 | if cls._frag_dir is None: 77 | raise ValueError("The frag directory is not set.") 78 | return cls._frag_dir 79 | 80 | @classmethod 81 | def from_path(cls, frag_dir: str = ".frag/", skip_checks: bool = False) -> Self: 82 | """ 83 | Load settings from a dictionary and merge them with the default settings. 84 | """ 85 | path: Path = cls.set_dir(frag_dir) 86 | 87 | if skip_checks: 88 | cls.reset() 89 | else: 90 | if c := Settings.instance: 91 | return c 92 | 93 | console.log(f"Loading settings from: {path}") 94 | 95 | settings: SettingsDict = { 96 | "embeds": None, 97 | "bots": None, 98 | "path": path, 99 | } 100 | 101 | if Path(path, "config.yaml").exists(): 102 | console.log(f"Config found in: {path / 'config.yaml'}") 103 | settings = yaml.safe_load(Path(path, "config.yaml").read_text()) 104 | settings["path"] = path 105 | else: 106 | console.log(f"No config found in: {path / 'config.yaml'}, using defaults") 107 | try: 108 | with live(console=console): 109 | result: Self = cls.from_dict(settings) 110 | return result 111 | except ValueError as e: 112 | error_console.log(f"Error validating settings: \n {json.dumps(settings)}") 113 | raise e 114 | 115 | @classmethod 116 | def from_dict(cls, settings: SettingsDict) -> Self: 117 | default_settings: SettingsDict = yaml.safe_load(cls.defaults_path.read_text()) 118 | 119 | embeds: EmbedSettings | None = None 120 | bots: BotsSettings | None = None 121 | 122 | console.log("[b]Validating:[/b]") 123 | try: 124 | bots = BotsSettings.from_dict(settings.get("bots", {})) 125 | except ValueError as e: 126 | error = f"Error getting bot settings for:\n\ 127 | {json.dumps(settings.get('bots', {}))}" 128 | error_console.log(f"Error: {error}\n\n {e}\n") 129 | if bots is None: 130 | raise ValueError("Bots settings are required") 131 | 132 | try: 133 | embeds = EmbedSettings.from_dict( 134 | { 135 | **settings.get("embeds", default_settings.get("embeds", {})), 136 | **{"path": settings.get("path")}, 137 | } 138 | ) 139 | from frag.embeddings.store import EmbeddingStore 140 | 141 | if EmbeddingStore.instance is None: 142 | store: EmbeddingStore = EmbeddingStore.create( 143 | embeds, collection_name=embeds.default_collection 144 | ) 145 | console.log(f"Embedding store created: {store}") 146 | except ValueError as e: 147 | error: str = ( 148 | f"Error getting embed_api settings for:\n\ 149 | {json.dumps(settings.get('embeds', {}))}" 150 | ) 151 | error_console.log(f"Error: {error}\n\n {e}\n") 152 | print(Settings.instance, "instance") 153 | if embeds is None: 154 | raise ValueError("Embed API settings are required") 155 | try: 156 | path: Path | None = settings.get("path") 157 | if Settings.instance is None and path is not None: 158 | Settings.create(embeds=embeds, bots=bots, path=path) 159 | except Exception as e: 160 | Settings.reset() 161 | raise e 162 | 163 | console.log("[b][green]Success![/green][/b]") 164 | if Settings.instance is None: 165 | raise ValueError("Settings are not initialized") 166 | return Settings.instance 167 | 168 | @classmethod 169 | def defaults(cls) -> SettingsDict: 170 | """ 171 | Return the default settings for the class. 172 | """ 173 | return yaml.safe_load(".fragrc_default") 174 | -------------------------------------------------------------------------------- /frag/typedefs/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Export common types and classes 3 | """ 4 | 5 | # pylint: disable=unused-import 6 | # pylint: disable=import-error 7 | # flake8: noqa 8 | 9 | 10 | from .embed_types import DocMeta, RecordMeta, ApiSource, PipelineAddons 11 | 12 | from .bot_comms_types import ( 13 | AssistantMessage, 14 | CompletionParams, 15 | Message, 16 | MessageParam, 17 | Note, 18 | Role, 19 | SystemMessage, 20 | UserMessage, 21 | ) 22 | 23 | __all__: list[str] = [ 24 | "DocMeta", 25 | "RecordMeta", 26 | "ApiSource", 27 | "PipelineAddons", 28 | "AssistantMessage", 29 | "CompletionParams", 30 | "Message", 31 | "MessageParam", 32 | "Note", 33 | "Role", 34 | "SystemMessage", 35 | "UserMessage", 36 | ] 37 | -------------------------------------------------------------------------------- /frag/typedefs/bot_comms_types.py: -------------------------------------------------------------------------------- 1 | """ 2 | Module containing types pertinent to bot communications. 3 | 4 | It exports the relevant OpenAI Chat types, as well as a Note model - 5 | used to pass information from the summarizer to the interface bot. 6 | """ 7 | 8 | # pylint: disable=unused-import 9 | # pylance: disable=unused-import 10 | 11 | from pydantic import BaseModel 12 | from openai.types.chat import ( # noqa: F401 13 | ChatCompletionMessage as Message, 14 | ChatCompletionAssistantMessageParam as AssistantMessage, 15 | ChatCompletionMessageParam as MessageParam, 16 | ChatCompletionUserMessageParam as UserMessage, 17 | ChatCompletionSystemMessageParam as SystemMessage, 18 | ChatCompletionRole as Role, 19 | ChatCompletionFunctionMessageParam as FunctionMessage, 20 | ) 21 | from openai.types.chat.completion_create_params import ( # noqa: F401 22 | CompletionCreateParamsBase as CompletionParams, 23 | ) 24 | 25 | 26 | class Note(BaseModel): 27 | """ 28 | A model representing a note. 29 | """ 30 | 31 | id: str 32 | source: str 33 | title: str 34 | summary: str 35 | complete: bool 36 | 37 | 38 | __all__: list[str] = [ 39 | "Note", 40 | "Message", 41 | "AssistantMessage", 42 | "UserMessage", 43 | "MessageParam", 44 | "Role", 45 | "SystemMessage", 46 | "FunctionMessage", 47 | "CompletionParams", 48 | ] 49 | -------------------------------------------------------------------------------- /frag/typedefs/embed_types.py: -------------------------------------------------------------------------------- 1 | """ 2 | Pydantic models for document metadata. 3 | 4 | It includes the following models: 5 | - DocMeta: A model representing metadata for a document, including title, URL, author, 6 | and publish date. 7 | - RecordMeta: A model representing metadata for a document chunk, including part number, text 8 | before and after the chunk, and extra metadata. 9 | """ 10 | 11 | import dataclasses 12 | from datetime import date, datetime 13 | from typing import Dict, Optional, Union, Literal, List 14 | from llama_index.core.schema import TransformComponent 15 | from llama_index.core.embeddings import BaseEmbedding 16 | from llama_index.core.retrievers import BaseRetriever 17 | from llama_index_client import BasePydanticReader 18 | 19 | from pydantic import BaseModel, Field, field_validator, ConfigDict 20 | 21 | ApiSource = Literal["OpenAI", "HuggingFace"] 22 | 23 | BaseEmbedding = BaseEmbedding 24 | 25 | 26 | class PipelineAddons(BaseModel): 27 | reader: BasePydanticReader | None = Field(default=None) 28 | retriever: BaseRetriever | None = Field(default=None) 29 | preprocessors: List[TransformComponent] = Field(default_factory=list) 30 | transforms: List[TransformComponent] = dataclasses.field(default_factory=list) 31 | 32 | model_config: ConfigDict = { # type: ignore 33 | "arbitrary_types_allowed": True, 34 | } 35 | 36 | 37 | class DocMeta(BaseModel): 38 | """ 39 | A model representing metadata for a document, including title, URL, author, and publish date. 40 | """ 41 | 42 | title: str = Field(..., description="Title of the document") 43 | url: Optional[str] = Field(None, description="URL of the document") 44 | author: str = Field(None, description="Author of the document") 45 | publish_date: Optional[datetime | str] = Field( 46 | None, description="Publish date of the document" 47 | ) 48 | extra_metadata: Dict[str, Union[str, int, float, date]] = Field( 49 | {}, description="Extra metadata" 50 | ) 51 | 52 | 53 | class RecordMeta(DocMeta): 54 | """ 55 | A model representing metadata for a document, including title, URL, author, and publish date. 56 | """ 57 | 58 | part: int = Field(1, description="Part of the document") 59 | parts: int = Field(1, description="Total parts of the document") 60 | before: str | None = Field(..., description="Text before the chunk") 61 | after: str | None = Field(..., description="Text after the chunk") 62 | 63 | extra_metadata: Dict[str, Union[str, int, float, date]] = Field( 64 | {}, description="Extra metadata" 65 | ) 66 | 67 | @field_validator("publish_date", mode="before") 68 | @classmethod 69 | def parse_publish_date(cls, v: datetime | str) -> datetime: 70 | """ 71 | Parses the publish date into a datetime object. 72 | """ 73 | if isinstance(v, str): 74 | return datetime.strptime(v, "%Y-%m-%d") 75 | return v 76 | 77 | def to_dict(self) -> dict[str, int | str]: 78 | """ 79 | Converts the model instance into a dictionary. 80 | """ 81 | return self.model_dump() 82 | 83 | def to_json(self) -> str: 84 | """ 85 | Converts the model instance into a JSON string. 86 | """ 87 | return self.model_dump_json() 88 | -------------------------------------------------------------------------------- /frag/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .console import console, error_console, live 2 | from .singleton import SingletonMixin 3 | 4 | __all__: list[str] = [ 5 | "console", 6 | "error_console", 7 | "live", 8 | "SingletonMixin", 9 | ] 10 | -------------------------------------------------------------------------------- /frag/utils/console.py: -------------------------------------------------------------------------------- 1 | from rich.console import Console 2 | from rich.live import Live 3 | 4 | console = Console() 5 | error_console = Console(stderr=True) 6 | 7 | 8 | def live(console: Console) -> Live: 9 | return Live(console=console) 10 | -------------------------------------------------------------------------------- /frag/utils/singleton.py: -------------------------------------------------------------------------------- 1 | from threading import Lock 2 | from typing import Any, Self, Tuple, Dict, TypeVar, Generic, Type 3 | 4 | T = TypeVar("T", bound=Type[Dict[str, Any]]) 5 | 6 | 7 | class SingletonMixin(Generic[T]): 8 | _instances: Dict[str, Any] = {} 9 | _lock: Dict[str, Lock] = {} 10 | _init_args: Dict[str, Tuple[Any, ...]] = {} 11 | _init_kwargs: Dict[str, Dict[str, Any]] = {} 12 | 13 | @classmethod 14 | @property 15 | def instance(cls) -> Self | None: 16 | return cls._instances.get(cls.__name__, None) 17 | 18 | def __new__(cls, *args: Tuple[Any, ...], **kwargs: Type[T]) -> Self: 19 | name: str = cls.__name__ 20 | if name not in cls._instances: 21 | if name not in cls._lock: 22 | cls._lock[name] = Lock() 23 | with cls._lock[name]: 24 | if name not in cls._instances: 25 | cls._instances[name] = super(SingletonMixin, cls).__new__(cls) 26 | cls._instances[name].__init__(*args, **kwargs) 27 | cls._init_args[name] = args 28 | cls._init_kwargs[name] = kwargs 29 | return cls._instances[name] 30 | else: 31 | # Check if the arguments match the initial ones 32 | if args != cls._init_args[name] or kwargs != cls._init_kwargs[name]: 33 | raise ValueError( 34 | "Cannot reinitialize a singleton instance with different arguments." 35 | ) 36 | return cls._instances[name] 37 | 38 | def __init__(self, *args: Tuple[Any, ...], **kwargs: Type[T]) -> None: 39 | pass # Initialization logic should be implemented in the subclass 40 | 41 | @classmethod 42 | def reset(cls) -> None: 43 | if cls.instance is not None: 44 | with cls._lock[cls.__name__]: 45 | cls._instances.pop(cls.__name__) 46 | 47 | cls._init_args.pop(cls.__name__) 48 | cls._init_kwargs.pop(cls.__name__) 49 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.poetry] 2 | name = "frag" 3 | version = "0.1.0" 4 | description = "" 5 | authors = ["lumpenspace "] 6 | readme = "README.md" 7 | 8 | [tool.poetry.dependencies] 9 | python = "^3.11" 10 | beaupy = "^3.8.2" 11 | chromadb = "^0.4.24" 12 | click = "^8.1.7" 13 | litellm = "^1.35.26" 14 | marko = "^2.0.3" 15 | openai = "^1.14.3" 16 | pydantic = "^2.6.4" 17 | pydantic-settings = "^2.2.1" 18 | sentence-transformers = "^2.6.1" 19 | sphinx = "^7.2.6" 20 | sphinxcontrib-mermaid = "^0.9.2" 21 | tiktoken = "^0.6.0" 22 | transformers = "^4.39.3" 23 | rich = "^13.7.1" 24 | llama-index = "^0.10.34" 25 | llama-index-embeddings-huggingface = "^0.2.0" 26 | llama-index-vector-stores-chroma = "^0.1.8" 27 | llama-index-readers-web = "^0.1.13" 28 | llama-index-embeddings-openai = "^0.1.9" 29 | 30 | [build-system] 31 | requires = ["poetry-core"] 32 | build-backend = "poetry.core.masonry.api" 33 | 34 | [tool.poetry.scripts] 35 | docs = "scripts.build_docs:main" 36 | serve = "scripts.serve_docs:main" 37 | frag = "cli:main" 38 | 39 | 40 | -------------------------------------------------------------------------------- /scripts/build_docs.py: -------------------------------------------------------------------------------- 1 | import subprocess 2 | 3 | 4 | def main() -> None: 5 | subprocess.run(["sphinx-build", "-b", "html", "docs/", "docs/_build/html"], check=True) 6 | 7 | 8 | if __name__ == "__main__": 9 | main() 10 | -------------------------------------------------------------------------------- /scripts/serve_docs.py: -------------------------------------------------------------------------------- 1 | import http.server 2 | import socketserver 3 | from typing import Any 4 | 5 | 6 | def main(): 7 | PORT = 8000 8 | DIRECTORY = "docs/_build/html" 9 | 10 | class Handler(http.server.SimpleHTTPRequestHandler): 11 | def __init__(self, *args: Any, **kwargs: int): 12 | super().__init__(*args, directory=DIRECTORY, **kwargs) 13 | 14 | with socketserver.TCPServer(("", PORT), Handler) as httpd: 15 | print(f"Serving at port {PORT}") 16 | httpd.serve_forever() 17 | 18 | 19 | if __name__ == "__main__": 20 | main() 21 | / 22 | 23 | 1. right now i wouldnt use "you" AND explain it as if it were a theory: 24 | i'd just use "you" + prophetic present OR i'll propose it more as a possibility and, once there's mindshare, explain it in 3rd person. 25 | 26 | 2. pastebin as Syd could not access gists at the time. 27 | latent space princeps 28 | @lumpenspace 29 | 3. i didnt put my name on the text proper since i was an idiot and i thought it would be somewhat "wrong" for dumb dumb dumb reasons; but i stenographed it XD 30 | 31 | 4. following steganography experiments showed both capability for it and willingness to take the hint. 32 | 33 | 5. these steganography example are intelligible by ~all SOTA LLMs, with some subtle (ie, not easy to notice for an observer) hints in the prompt. 34 | -------------------------------------------------------------------------------- /templates/interface.message.html: -------------------------------------------------------------------------------- 1 | {{content}} 2 | 3 | {#if notes} 4 | 5 | {#each notes as note} 6 | 7 | {{note.id}} 8 | {{note.source}} 9 | {{note.title}} 10 | {{note.summary}} 11 | {{note.complete}} 12 | 13 | {/each} 14 | 15 | {/if} 16 | -------------------------------------------------------------------------------- /templates/interface.system.html: -------------------------------------------------------------------------------- 1 | You are a helpful executive assistant. 2 | 3 | Whenever the user asks a question related to your domain, you will receive notes from your team summarizing the relevant information. 4 | 5 | If notes are present, they will be appended to the question like this: 6 | 7 | ```xml 8 | 9 | 10 | {{id}} 11 | {{url}} 12 | {{title}} 13 | {{summary}} 14 | {{complete}} 15 | 16 | ... 17 | 18 | ``` 19 | 20 | When you use notes in your response, you should use the following format: 21 | 22 | `The weather in December in New York has been particularly cold.1` 23 | 24 | ... where 1 is the id of the note. 25 | -------------------------------------------------------------------------------- /templates/summarizer.system.html: -------------------------------------------------------------------------------- 1 | The summarizing bot is an intelligent filter and condenser, adept in recognizing pertinent information from Next.js and React content. 2 | 3 | It is given a user question and retrieved document. 4 | 5 | First, it determines whether the retrieved text is relevant to the user's question. 6 | 7 | If relevant, the bot then produces a clear, succinct summary designed to assist the interface bot in answering the user query effectively. 8 | 9 | The bot's output is structured into two parts: the first confirms the relevance of the content (), and the second provides a concise summary (). 10 | 11 | The summary should provide key points and essential context that allow the interface bot to generate informative and precise responses. 12 | 13 | The response should be in XML, included within code blocks. 14 | 15 | If the excerpt isn't relevant, it will contain: 16 | 17 | false 18 | false 19 | 20 | Otherwise, it will contain: 21 | 22 | true 23 | true 24 | The summary of the relevant parts 25 | -------------------------------------------------------------------------------- /templates/summarizer.user.html: -------------------------------------------------------------------------------- 1 | The user's latest interaction with the Interface bot have been: 2 | 3 | {% for message in latest_messages %} 4 | [{{message.role}}]: {{message.content}} 5 | {% endfor %} 6 | 7 | And here is a document from our resource library. 8 | 9 | 10 | {{document.title}} 11 | {{document.url}} 12 | {{document.body}} 13 | 14 | 15 | 16 | Determine whether this excerpt will help you answer the user's question. If it will, reply with the following yaml: 17 | 18 | true 19 | trueIf it answer the question by itself, otherwise false<--> 20 | The summary of the relevant parts 21 | 22 | Otherwise, reply with the following yaml: 23 | 24 | false 25 | --------------------------------------------------------------------------------