diff --git a/examples/local_llm/local_llm.ipynb b/examples/local_llm/local_llm.ipynb index 6ce93bd2..17e97571 100644 --- a/examples/local_llm/local_llm.ipynb +++ b/examples/local_llm/local_llm.ipynb @@ -127,7 +127,7 @@ "output_type": "stream", "text": [ "User: What is Ragna?\n", - "Assistant: Ragna is an open-source application (OSS) for RAG workflows. It offers a Python and REST API as well as a web UI.\n" + "Assistant: Ragna is an OSS app for RAG workflows that offers a Python and REST API as well as web UI.\n" ] } ], @@ -143,13 +143,14 @@ " \"Ragna is an OSS app for RAG workflows that offers a Python and REST API as well as web UI\\n\"\n", " )\n", "\n", - "async with await rag.new_chat(\n", + "async with rag.chat(\n", " documents=[path],\n", " source_storage=RagnaDemoSourceStorage,\n", " assistant=AiroborosAssistant,\n", ") as chat:\n", " prompt = \"What is Ragna?\"\n", - " answer = await chat.answer(prompt)\n", + " message = await chat.answer(prompt)\n", + " answer = message.content\n", "\n", "print(f\"User: {prompt}\")\n", "print(f\"Assistant: {answer}\")" diff --git a/examples/python_api/python_api.ipynb b/examples/python_api/python_api.ipynb index 9380aad3..65489cc1 100644 --- a/examples/python_api/python_api.ipynb +++ b/examples/python_api/python_api.ipynb @@ -44,7 +44,7 @@ { "data": { "text/plain": [ - "Config(local_cache_root=PosixPath('/home/philip/.cache/ragna'), state_database_url='sqlite://', queue_database_url='memory', ragna_api_url='http://127.0.0.1:31476', ragna_ui_url='http://127.0.0.1:31477', document_class=, upload_token_secret='a9e79b896f90c136820e70ced28f30e6cbbaca6d19f52e41c35b154df21a657a', upload_token_ttl=30, registered_source_storage_classes={'Ragna/DemoSourceStorage': }, registered_assistant_classes={'Ragna/DemoAssistant': })" + "Config(local_cache_root=PosixPath('/home/philip/.cache/ragna'), state_database_url='sqlite://', queue_database_url='memory', ragna_api_url='http://127.0.0.1:31476', ragna_ui_url='http://127.0.0.1:31477', document_class=, upload_token_secret='9ffa4b72f96f221953a455df1642304af954a0afa7d147ae96cc1aa62a89ed02', upload_token_ttl=30, registered_source_storage_classes={'Ragna/DemoSourceStorage': }, registered_assistant_classes={'Ragna/DemoAssistant': })" ] }, "execution_count": 3, @@ -106,7 +106,7 @@ "\n", "rag = Rag(demo_config)\n", "\n", - "async with await rag.new_chat(\n", + "async with rag.chat(\n", " documents=[document_path],\n", " source_storage=RagnaDemoSourceStorage,\n", " assistant=RagnaDemoAssistant,\n", @@ -127,10 +127,10 @@ "name": "stdout", "output_type": "stream", "text": [ - "{('Chroma', 'OpenAI/gpt-3.5-turbo-16k'): ,\n", - " ('Chroma', 'OpenAI/gpt-4'): ,\n", - " ('LanceDB', 'OpenAI/gpt-3.5-turbo-16k'): ,\n", - " ('LanceDB', 'OpenAI/gpt-4'): }\n" + "{('Chroma', 'OpenAI/gpt-3.5-turbo-16k'): ,\n", + " ('Chroma', 'OpenAI/gpt-4'): ,\n", + " ('LanceDB', 'OpenAI/gpt-3.5-turbo-16k'): ,\n", + " ('LanceDB', 'OpenAI/gpt-4'): }\n" ] } ], @@ -144,13 +144,13 @@ "\n", "\n", "async def answer_prompt(source_storage, assistant):\n", - " chat = await rag.new_chat(\n", + " async with rag.chat(\n", " documents=[document_path],\n", " source_storage=source_storage,\n", " assistant=assistant,\n", - " )\n", - " await chat.start()\n", - " return await chat.answer(prompt)\n", + " ) as chat:\n", + " message = await chat.answer(prompt)\n", + " return message.content\n", "\n", "\n", "experiments = {\n", @@ -177,18 +177,35 @@ "To disable this warning, you can either:\n", "\t- Avoid using `tokenizers` before the fork if possible\n", "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n", - "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n", - "To disable this warning, you can either:\n", - "\t- Avoid using `tokenizers` before the fork if possible\n", - "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n", - "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n", - "To disable this warning, you can either:\n", - "\t- Avoid using `tokenizers` before the fork if possible\n", - "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n", - "{('Chroma', 'OpenAI/gpt-3.5-turbo-16k'): Ragna is an open-source RAG (Response Analysis Graph) orchestration app. It is designed to help users create conversational AI applications by providing a framework for managing and orchestrating the flow of conversations. Ragna allows developers to define conversation flows, handle user inputs, and generate dynamic responses based on predefined rules and logic. It is built on top of the Rasa framework and provides additional features and functionalities to simplify the development process.,\n", - " ('Chroma', 'OpenAI/gpt-4'): Ragna is an open-source RAG orchestration app.,\n", - " ('LanceDB', 'OpenAI/gpt-3.5-turbo-16k'): Ragna is an open-source rag orchestration app. It is a software application that allows users to create and arrange musical compositions using ragtime music. It is designed to be accessible and customizable for musicians and composers.,\n", - " ('LanceDB', 'OpenAI/gpt-4'): Ragna is an open-source rag orchestration app.}\n" + "{('Chroma', 'OpenAI/gpt-3.5-turbo-16k'): 'Ragna is an open-source RAG '\n", + " '(Response Analysis Graph) '\n", + " 'orchestration app. It is designed to '\n", + " 'help users create conversational AI '\n", + " 'applications by providing a '\n", + " 'framework for managing and '\n", + " 'orchestrating the flow of '\n", + " 'conversations. Ragna allows '\n", + " 'developers to define conversation '\n", + " 'flows, handle user inputs, and '\n", + " 'generate dynamic responses based on '\n", + " 'predefined rules and logic. It is '\n", + " 'built on top of the Rasa framework '\n", + " 'and provides additional features and '\n", + " 'functionalities to simplify the '\n", + " 'development process.',\n", + " ('Chroma', 'OpenAI/gpt-4'): 'Ragna is an open-source RAG orchestration app. '\n", + " 'Unfortunately, without additional context or '\n", + " \"sources, I can't provide more detailed \"\n", + " 'information about it.',\n", + " ('LanceDB', 'OpenAI/gpt-3.5-turbo-16k'): 'Ragna is an open-source rag '\n", + " 'orchestration app. It is a software '\n", + " 'application that allows users to '\n", + " 'create and arrange musical '\n", + " 'compositions using ragtime music. '\n", + " 'It is designed to be accessible and '\n", + " 'customizable for musicians and '\n", + " 'composers.',\n", + " ('LanceDB', 'OpenAI/gpt-4'): 'Ragna is an open-source rag orchestration app.'}\n" ] } ], diff --git a/examples/rest_api/rest_api.ipynb b/examples/rest_api/rest_api.ipynb index 7dff2d34..134646fd 100644 --- a/examples/rest_api/rest_api.ipynb +++ b/examples/rest_api/rest_api.ipynb @@ -41,20 +41,20 @@ "metadata": {}, "outputs": [ { - "name": "stdout", + "name": "stderr", "output_type": "stream", "text": [ - "INFO: 127.0.0.1:60998 - \"GET /health HTTP/1.1\" 200 OK\n" + "INFO: Started server process [32871]\n", + "INFO: Waiting for application startup.\n", + "INFO: Application startup complete.\n", + "INFO: Uvicorn running on http://127.0.0.1:31476 (Press CTRL+C to quit)\n" ] }, { - "name": "stderr", + "name": "stdout", "output_type": "stream", "text": [ - "INFO: Started server process [32153]\n", - "INFO: Waiting for application startup.\n", - "INFO: Application startup complete.\n", - "INFO: Uvicorn running on http://127.0.0.1:31476 (Press CTRL+C to quit)\n" + "INFO: 127.0.0.1:43564 - \"GET / HTTP/1.1\" 200 OK\n" ] } ], @@ -73,7 +73,7 @@ "start = time.time()\n", "while (time.time() - start) < timeout:\n", " with contextlib.suppress(httpx.ConnectError):\n", - " response = await client.get(f\"{URL}/health\")\n", + " response = await client.get(URL)\n", " if response.is_success:\n", " break\n", "\n", @@ -137,7 +137,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "INFO: 127.0.0.1:60998 - \"GET /chats?user=Ragna HTTP/1.1\" 200 OK\n", + "INFO: 127.0.0.1:43564 - \"GET /chats?user=Ragna HTTP/1.1\" 200 OK\n", "[]\n" ] } @@ -167,7 +167,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "INFO: 127.0.0.1:60998 - \"GET /components?user=Ragna HTTP/1.1\" 200 OK\n", + "INFO: 127.0.0.1:43564 - \"GET /components?user=Ragna HTTP/1.1\" 200 OK\n", "{'assistants': ['Ragna/DemoAssistant'],\n", " 'source_storages': ['Ragna/DemoSourceStorage']}\n" ] @@ -229,9 +229,9 @@ "name": "stdout", "output_type": "stream", "text": [ - "INFO: 127.0.0.1:60998 - \"GET /document?user=Ragna&name=document0.txt HTTP/1.1\" 200 OK\n", - "{'data': {'token': 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VyIjoiUmFnbmEiLCJpZCI6ImRlNWJiYjMxLTQ1MDEtNDFhNC05ZDRlLTFiN2EzNjk1Mjk5NSIsImV4cCI6MTY5Njg3Nzk1OS40NTE3NzcyfQ.XC7qKuCZXkbrI7isQ39IqYZMCL_O6oVP4GfI8wqHpdE'},\n", - " 'document': {'id': 'de5bbb31-4501-41a4-9d4e-1b7a36952995',\n", + "INFO: 127.0.0.1:43564 - \"GET /document?user=Ragna&name=document0.txt HTTP/1.1\" 200 OK\n", + "{'data': {'token': 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VyIjoiUmFnbmEiLCJpZCI6ImZlZTE4ZDc1LTE3NWQtNDdjOS05MTZhLTM1MTIxNzhhNDUwYiIsImV4cCI6MTY5NzExNDcwOS45OTgwOTU1fQ.3i1s_xc-nRFTuj8hlczipsZk0F7X1Vo5A9kmSau-guI'},\n", + " 'document': {'id': 'fee18d75-175d-47c9-916a-3512178a450b',\n", " 'name': 'document0.txt'},\n", " 'url': 'http://127.0.0.1:31476/document'}\n" ] @@ -264,7 +264,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "INFO: 127.0.0.1:60998 - \"POST /document HTTP/1.1\" 200 OK\n" + "INFO: 127.0.0.1:43564 - \"POST /document HTTP/1.1\" 200 OK\n" ] } ], @@ -295,18 +295,18 @@ "name": "stdout", "output_type": "stream", "text": [ - "INFO: 127.0.0.1:60998 - \"GET /document?user=Ragna&name=document1.txt HTTP/1.1\" 200 OK\n", - "INFO: 127.0.0.1:60998 - \"POST /document HTTP/1.1\" 200 OK\n", - "INFO: 127.0.0.1:60998 - \"GET /document?user=Ragna&name=document2.txt HTTP/1.1\" 200 OK\n", - "INFO: 127.0.0.1:60998 - \"POST /document HTTP/1.1\" 200 OK\n" + "INFO: 127.0.0.1:43564 - \"GET /document?user=Ragna&name=document1.txt HTTP/1.1\" 200 OK\n", + "INFO: 127.0.0.1:43564 - \"POST /document HTTP/1.1\" 200 OK\n", + "INFO: 127.0.0.1:43564 - \"GET /document?user=Ragna&name=document2.txt HTTP/1.1\" 200 OK\n", + "INFO: 127.0.0.1:43564 - \"POST /document HTTP/1.1\" 200 OK\n" ] }, { "data": { "text/plain": [ - "[{'id': 'de5bbb31-4501-41a4-9d4e-1b7a36952995', 'name': 'document0.txt'},\n", - " {'id': '8e5e17b4-e6e3-4707-806a-917f0f5b46fb', 'name': 'document1.txt'},\n", - " {'id': '336436fa-1b6a-4a29-8821-791cd682a7cb', 'name': 'document2.txt'}]" + "[{'id': 'fee18d75-175d-47c9-916a-3512178a450b', 'name': 'document0.txt'},\n", + " {'id': '8c73c112-45a9-457d-ada8-87b3e63e5b9a', 'name': 'document1.txt'},\n", + " {'id': 'd62a61ec-c877-403f-8723-19b647397edc', 'name': 'document2.txt'}]" ] }, "execution_count": 9, @@ -349,21 +349,20 @@ "name": "stdout", "output_type": "stream", "text": [ - "INFO: 127.0.0.1:60998 - \"POST /chats?user=Ragna HTTP/1.1\" 200 OK\n", - "{'closed': False,\n", - " 'id': '0777396d-9993-47ce-b807-2e9c48364c63',\n", - " 'messages': [],\n", - " 'metadata': {'assistant': 'Ragna/DemoAssistant',\n", - " 'documents': [{'id': 'de5bbb31-4501-41a4-9d4e-1b7a36952995',\n", + "INFO: 127.0.0.1:43564 - \"POST /chats?user=Ragna HTTP/1.1\" 200 OK\n", + "{'id': '9d2e7a49-3030-4ea8-a2a2-1e0b7dad7f86',\n", + " 'metadata': {'name': 'Ragna REST API example',\n", + " 'source_storage': 'Ragna/DemoSourceStorage',\n", + " 'assistant': 'Ragna/DemoAssistant',\n", + " 'params': {},\n", + " 'documents': [{'id': 'fee18d75-175d-47c9-916a-3512178a450b',\n", " 'name': 'document0.txt'},\n", - " {'id': '8e5e17b4-e6e3-4707-806a-917f0f5b46fb',\n", + " {'id': '8c73c112-45a9-457d-ada8-87b3e63e5b9a',\n", " 'name': 'document1.txt'},\n", - " {'id': '336436fa-1b6a-4a29-8821-791cd682a7cb',\n", - " 'name': 'document2.txt'}],\n", - " 'name': 'Ragna REST API example',\n", - " 'params': {},\n", - " 'source_storage': 'Ragna/DemoSourceStorage'},\n", - " 'started': False}\n" + " {'id': 'd62a61ec-c877-403f-8723-19b647397edc',\n", + " 'name': 'document2.txt'}]},\n", + " 'messages': [],\n", + " 'prepared': False}\n" ] } ], @@ -373,14 +372,14 @@ " params={\"user\": USER},\n", " json={\n", " \"name\": \"Ragna REST API example\",\n", - " \"document_ids\": [d[\"id\"] for d in documents],\n", + " \"documents\": documents,\n", " \"source_storage\": SOURCE_STORAGE,\n", " \"assistant\": ASSISTANT,\n", " \"params\": {},\n", " },\n", ")\n", "chat = response.json()\n", - "pprint(chat)" + "pprint(chat, sort_dicts=False)" ] }, { @@ -388,7 +387,7 @@ "id": "d7b74f96-3e02-4a70-8d54-c34d73f706bd", "metadata": {}, "source": [ - "As indicated by the `'started': False` in the response, we need to start our chat before we can start the interogation. In this step we extract the data out of the uploaded documents and store them in our source storage" + "As indicated by the `'prepared': False` in the response, we need to prepare our chat before we can start the interogation. In this step we extract the data out of the uploaded documents and store them in our source storage" ] }, { @@ -401,34 +400,38 @@ "name": "stdout", "output_type": "stream", "text": [ - "INFO: 127.0.0.1:60998 - \"POST /chats/0777396d-9993-47ce-b807-2e9c48364c63/start?user=Ragna HTTP/1.1\" 200 OK\n", - "{'closed': False,\n", - " 'id': '0777396d-9993-47ce-b807-2e9c48364c63',\n", - " 'messages': [{'content': 'How can I help you with the documents?',\n", - " 'id': 'f9b43a27-67ac-4984-9deb-3844ff4a6ea1',\n", - " 'role': 'system',\n", - " 'sources': [],\n", - " 'timestamp': '2023-10-09T18:58:50.356957'}],\n", - " 'metadata': {'assistant': 'Ragna/DemoAssistant',\n", - " 'documents': [{'id': '336436fa-1b6a-4a29-8821-791cd682a7cb',\n", - " 'name': 'document2.txt'},\n", - " {'id': '8e5e17b4-e6e3-4707-806a-917f0f5b46fb',\n", - " 'name': 'document1.txt'},\n", - " {'id': 'de5bbb31-4501-41a4-9d4e-1b7a36952995',\n", - " 'name': 'document0.txt'}],\n", - " 'name': 'Ragna REST API example',\n", - " 'params': {},\n", - " 'source_storage': 'Ragna/DemoSourceStorage'},\n", - " 'started': True}\n" + "INFO: 127.0.0.1:43564 - \"POST /chats/9d2e7a49-3030-4ea8-a2a2-1e0b7dad7f86/prepare?user=Ragna HTTP/1.1\" 200 OK\n", + "{'message': {'id': 'a1a3b272-7145-4459-b89c-1de258bb7cff',\n", + " 'content': 'How can I help you with the documents?',\n", + " 'role': 'system',\n", + " 'sources': [],\n", + " 'timestamp': '2023-10-12T12:44:40.986403+00:00'},\n", + " 'chat': {'id': '9d2e7a49-3030-4ea8-a2a2-1e0b7dad7f86',\n", + " 'metadata': {'name': 'Ragna REST API example',\n", + " 'source_storage': 'Ragna/DemoSourceStorage',\n", + " 'assistant': 'Ragna/DemoAssistant',\n", + " 'params': {},\n", + " 'documents': [{'id': '8c73c112-45a9-457d-ada8-87b3e63e5b9a',\n", + " 'name': 'document1.txt'},\n", + " {'id': 'd62a61ec-c877-403f-8723-19b647397edc',\n", + " 'name': 'document2.txt'},\n", + " {'id': 'fee18d75-175d-47c9-916a-3512178a450b',\n", + " 'name': 'document0.txt'}]},\n", + " 'messages': [{'id': 'a1a3b272-7145-4459-b89c-1de258bb7cff',\n", + " 'content': 'How can I help you with the documents?',\n", + " 'role': 'system',\n", + " 'sources': [],\n", + " 'timestamp': '2023-10-12T12:44:40.986403+00:00'}],\n", + " 'prepared': True}}\n" ] } ], "source": [ "CHAT_ID = chat[\"id\"]\n", "\n", - "response = await client.post(f\"{URL}/chats/{CHAT_ID}/start\", params={\"user\": USER})\n", + "response = await client.post(f\"{URL}/chats/{CHAT_ID}/prepare\", params={\"user\": USER})\n", "chat = response.json()\n", - "pprint(chat)" + "pprint(chat, sort_dicts=False)" ] }, { @@ -449,49 +452,65 @@ "name": "stdout", "output_type": "stream", "text": [ - "INFO: 127.0.0.1:60998 - \"POST /chats/0777396d-9993-47ce-b807-2e9c48364c63/answer?user=Ragna&prompt=What%20is%20Ragna%3F HTTP/1.1\" 200 OK\n", - "{'content': \"I can't really help you with your prompt:\\n\"\n", + "INFO: 127.0.0.1:43564 - \"POST /chats/9d2e7a49-3030-4ea8-a2a2-1e0b7dad7f86/answer?user=Ragna&prompt=What%20is%20Ragna%3F HTTP/1.1\" 200 OK\n", + "{'id': 'b96faacd-e941-4d81-ae3c-0dfcbb1bf5df',\n", + " 'content': \"I can't really help you with your prompt:\\n\"\n", " '\\n'\n", " '> What is Ragna?\\n'\n", " '\\n'\n", " 'I can at least show you the sources that I was given:\\n'\n", " '\\n'\n", - " '- document2.txt: This is content of document 2\\n'\n", " '- document1.txt: This is content of document 1\\n'\n", + " '- document2.txt: This is content of document 2\\n'\n", " '- document0.txt: This is content of document 0',\n", - " 'id': '00f34df9-5e74-4588-8808-e1e11987db7e',\n", " 'role': 'assistant',\n", - " 'sources': [{'document': {'id': '336436fa-1b6a-4a29-8821-791cd682a7cb',\n", - " 'name': 'document2.txt'},\n", - " 'id': 'a07d9305-860b-45c7-9997-bdf47b7dfe58',\n", - " 'location': ''},\n", - " {'document': {'id': '8e5e17b4-e6e3-4707-806a-917f0f5b46fb',\n", + " 'sources': [{'id': 'c78a2488-bc95-4448-9c49-588ff94387e1',\n", + " 'document': {'id': '8c73c112-45a9-457d-ada8-87b3e63e5b9a',\n", " 'name': 'document1.txt'},\n", - " 'id': '1fac7064-f8c0-4a85-9e52-06ade4bfff86',\n", " 'location': ''},\n", - " {'document': {'id': 'de5bbb31-4501-41a4-9d4e-1b7a36952995',\n", + " {'id': '8784eae4-df01-43c1-bab0-3ac1ab7388b6',\n", + " 'document': {'id': 'd62a61ec-c877-403f-8723-19b647397edc',\n", + " 'name': 'document2.txt'},\n", + " 'location': ''},\n", + " {'id': 'aab43ed1-b6b8-4cd3-addc-238fa1db4308',\n", + " 'document': {'id': 'fee18d75-175d-47c9-916a-3512178a450b',\n", " 'name': 'document0.txt'},\n", - " 'id': '8e5e94d7-35ab-4426-b369-88e04ef19e80',\n", " 'location': ''}],\n", - " 'timestamp': '2023-10-09T18:58:50.368654'}\n", + " 'timestamp': '2023-10-12T12:44:41.000387+00:00'}\n" + ] + } + ], + "source": [ + "response = await client.post(\n", + " f\"{URL}/chats/{CHAT_ID}/answer\", params={\"user\": USER, \"prompt\": \"What is Ragna?\"}\n", + ")\n", + "answer = response.json()\n", + "pprint(answer[\"message\"], sort_dicts=False)" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "87f84fe0-079d-4c74-8995-4bf1749432ba", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ "I can't really help you with your prompt:\n", "\n", "> What is Ragna?\n", "\n", "I can at least show you the sources that I was given:\n", "\n", - "- document2.txt: This is content of document 2\n", "- document1.txt: This is content of document 1\n", + "- document2.txt: This is content of document 2\n", "- document0.txt: This is content of document 0\n" ] } ], "source": [ - "response = await client.post(\n", - " f\"{URL}/chats/{CHAT_ID}/answer\", params={\"user\": USER, \"prompt\": \"What is Ragna?\"}\n", - ")\n", - "answer = response.json()\n", - "pprint(answer[\"message\"])\n", "print(answer[\"message\"][\"content\"])" ] }, @@ -500,75 +519,64 @@ "id": "38071c7b-13a7-49c5-84f5-e8f1d205d578", "metadata": {}, "source": [ - "Welp, that was not really helpful, but unfortunately, this is the reality for the demo components we selected. Select some more elaborate components and you will get better answers. We could keep keep requesting answers, but at some point, the user likely wants to close the chat and move on. Doing so will prevent any further questions to be asked." + "Welp, that was not really helpful, but unfortunately, this is the reality for the demo components we selected. Select some more elaborate components and you will get better answers. We could keep keep requesting answers, but at some point, the user likely wants to delete the chat and move on." ] }, { "cell_type": "code", - "execution_count": 13, - "id": "f021fe9b-ae71-4101-aa18-22277ebab8d2", + "execution_count": 14, + "id": "1d5834d7-b269-4530-a0fe-67b3b3d7bda7", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ - "INFO: 127.0.0.1:60998 - \"POST /chats/0777396d-9993-47ce-b807-2e9c48364c63/close?user=Ragna HTTP/1.1\" 200 OK\n", - "{'closed': True,\n", - " 'id': '0777396d-9993-47ce-b807-2e9c48364c63',\n", - " 'messages': [{'content': 'How can I help you with the documents?',\n", - " 'id': 'f9b43a27-67ac-4984-9deb-3844ff4a6ea1',\n", - " 'role': 'system',\n", - " 'sources': [],\n", - " 'timestamp': '2023-10-09T18:58:50.356957'},\n", - " {'content': 'What is Ragna?',\n", - " 'id': '1d7f401a-556e-4c00-af11-8765c01d1709',\n", - " 'role': 'user',\n", - " 'sources': [],\n", - " 'timestamp': '2023-10-09T18:58:50.366988'},\n", - " {'content': \"I can't really help you with your prompt:\\n\"\n", - " '\\n'\n", - " '> What is Ragna?\\n'\n", - " '\\n'\n", - " 'I can at least show you the sources that I was '\n", - " 'given:\\n'\n", - " '\\n'\n", - " '- document2.txt: This is content of document 2\\n'\n", - " '- document1.txt: This is content of document 1\\n'\n", - " '- document0.txt: This is content of document 0',\n", - " 'id': '00f34df9-5e74-4588-8808-e1e11987db7e',\n", - " 'role': 'assistant',\n", - " 'sources': [{'document': {'id': '336436fa-1b6a-4a29-8821-791cd682a7cb',\n", - " 'name': 'document2.txt'},\n", - " 'id': 'a07d9305-860b-45c7-9997-bdf47b7dfe58',\n", - " 'location': ''},\n", - " {'document': {'id': '8e5e17b4-e6e3-4707-806a-917f0f5b46fb',\n", - " 'name': 'document1.txt'},\n", - " 'id': '1fac7064-f8c0-4a85-9e52-06ade4bfff86',\n", - " 'location': ''},\n", - " {'document': {'id': 'de5bbb31-4501-41a4-9d4e-1b7a36952995',\n", - " 'name': 'document0.txt'},\n", - " 'id': '8e5e94d7-35ab-4426-b369-88e04ef19e80',\n", - " 'location': ''}],\n", - " 'timestamp': '2023-10-09T18:58:50.368654'}],\n", - " 'metadata': {'assistant': 'Ragna/DemoAssistant',\n", - " 'documents': [{'id': '336436fa-1b6a-4a29-8821-791cd682a7cb',\n", - " 'name': 'document2.txt'},\n", - " {'id': '8e5e17b4-e6e3-4707-806a-917f0f5b46fb',\n", - " 'name': 'document1.txt'},\n", - " {'id': 'de5bbb31-4501-41a4-9d4e-1b7a36952995',\n", - " 'name': 'document0.txt'}],\n", - " 'name': 'Ragna REST API example',\n", - " 'params': {},\n", - " 'source_storage': 'Ragna/DemoSourceStorage'},\n", - " 'started': True}\n" + "INFO: 127.0.0.1:43564 - \"DELETE /chats/9d2e7a49-3030-4ea8-a2a2-1e0b7dad7f86?user=Ragna HTTP/1.1\" 200 OK\n" ] + }, + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 14, + "metadata": {}, + "output_type": "execute_result" } ], "source": [ - "response = await client.post(f\"{URL}/chats/{CHAT_ID}/close\", params={\"user\": USER})\n", - "chat = response.json()\n", - "pprint(chat)" + "await client.delete(f\"{URL}/chats/{CHAT_ID}\", params={\"user\": USER})" + ] + }, + { + "cell_type": "markdown", + "id": "f9445d81-8db5-41dd-b11b-ea18f77eb6b5", + "metadata": {}, + "source": [ + "After the chat is deleted, querying all chats of a user returns an empty list." + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "b1a027b3-be97-4364-a0a8-0c3785014f25", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "INFO: 127.0.0.1:43564 - \"GET /chats?user=Ragna HTTP/1.1\" 200 OK\n", + "[]\n" + ] + } + ], + "source": [ + "response = await client.get(f\"{URL}/chats\", params={\"user\": USER})\n", + "chats = response.json()\n", + "pprint(chats, sort_dicts=True)" ] } ], diff --git a/examples/s3_documents/s3_document_config.py b/examples/s3_documents/s3_document_config.py index 940f2305..cb6ca136 100644 --- a/examples/s3_documents/s3_document_config.py +++ b/examples/s3_documents/s3_document_config.py @@ -1,4 +1,6 @@ import os + +import uuid from typing import Any from ragna.assistant import RagnaDemoAssistant @@ -20,7 +22,7 @@ def _session(cls): @classmethod async def get_upload_info( - cls, *, config: Config, user: str, id: str, name: str + cls, *, config: Config, user: str, id: uuid.UUID, name: str ) -> tuple[str, dict[str, Any], dict[str, Any]]: if not PackageRequirement("boto3").is_available(): raise RagnaException() diff --git a/examples/s3_documents/s3_documents.ipynb b/examples/s3_documents/s3_documents.ipynb index f1a33a5d..d93f5ccb 100644 --- a/examples/s3_documents/s3_documents.ipynb +++ b/examples/s3_documents/s3_documents.ipynb @@ -135,6 +135,8 @@ ".output_html .vi { color: #19177C } /* Name.Variable.Instance */\n", ".output_html .vm { color: #19177C } /* Name.Variable.Magic */\n", ".output_html .il { color: #666666 } /* Literal.Number.Integer.Long */
import os\n",
+       "\n",
+       "import uuid\n",
        "from typing import Any\n",
        "\n",
        "from ragna.assistant import RagnaDemoAssistant\n",
@@ -156,7 +158,7 @@
        "\n",
        "    @classmethod\n",
        "    async def get_upload_info(\n",
-       "        cls, *, config: Config, user: str, id: str, name: str\n",
+       "        cls, *, config: Config, user: str, id: uuid.UUID, name: str\n",
        "    ) -> tuple[str, dict[str, Any], dict[str, Any]]:\n",
        "        if not PackageRequirement("boto3").is_available():\n",
        "            raise RagnaException()\n",
@@ -210,6 +212,8 @@
       "text/latex": [
        "\\begin{Verbatim}[commandchars=\\\\\\{\\}]\n",
        "\\PY{k+kn}{import} \\PY{n+nn}{os}\n",
+       "\n",
+       "\\PY{k+kn}{import} \\PY{n+nn}{uuid}\n",
        "\\PY{k+kn}{from} \\PY{n+nn}{typing} \\PY{k+kn}{import} \\PY{n}{Any}\n",
        "\n",
        "\\PY{k+kn}{from} \\PY{n+nn}{ragna}\\PY{n+nn}{.}\\PY{n+nn}{assistant} \\PY{k+kn}{import} \\PY{n}{RagnaDemoAssistant}\n",
@@ -231,7 +235,7 @@
        "\n",
        "    \\PY{n+nd}{@classmethod}\n",
        "    \\PY{k}{async} \\PY{k}{def} \\PY{n+nf}{get\\PYZus{}upload\\PYZus{}info}\\PY{p}{(}\n",
-       "        \\PY{n+nb+bp}{cls}\\PY{p}{,} \\PY{o}{*}\\PY{p}{,} \\PY{n}{config}\\PY{p}{:} \\PY{n}{Config}\\PY{p}{,} \\PY{n}{user}\\PY{p}{:} \\PY{n+nb}{str}\\PY{p}{,} \\PY{n+nb}{id}\\PY{p}{:} \\PY{n+nb}{str}\\PY{p}{,} \\PY{n}{name}\\PY{p}{:} \\PY{n+nb}{str}\n",
+       "        \\PY{n+nb+bp}{cls}\\PY{p}{,} \\PY{o}{*}\\PY{p}{,} \\PY{n}{config}\\PY{p}{:} \\PY{n}{Config}\\PY{p}{,} \\PY{n}{user}\\PY{p}{:} \\PY{n+nb}{str}\\PY{p}{,} \\PY{n+nb}{id}\\PY{p}{:} \\PY{n}{uuid}\\PY{o}{.}\\PY{n}{UUID}\\PY{p}{,} \\PY{n}{name}\\PY{p}{:} \\PY{n+nb}{str}\n",
        "    \\PY{p}{)} \\PY{o}{\\PYZhy{}}\\PY{o}{\\PYZgt{}} \\PY{n+nb}{tuple}\\PY{p}{[}\\PY{n+nb}{str}\\PY{p}{,} \\PY{n+nb}{dict}\\PY{p}{[}\\PY{n+nb}{str}\\PY{p}{,} \\PY{n}{Any}\\PY{p}{]}\\PY{p}{,} \\PY{n+nb}{dict}\\PY{p}{[}\\PY{n+nb}{str}\\PY{p}{,} \\PY{n}{Any}\\PY{p}{]}\\PY{p}{]}\\PY{p}{:}\n",
        "        \\PY{k}{if} \\PY{o+ow}{not} \\PY{n}{PackageRequirement}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{boto3}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{o}{.}\\PY{n}{is\\PYZus{}available}\\PY{p}{(}\\PY{p}{)}\\PY{p}{:}\n",
        "            \\PY{k}{raise} \\PY{n}{RagnaException}\\PY{p}{(}\\PY{p}{)}\n",
@@ -284,6 +288,8 @@
       ],
       "text/plain": [
        "import os\n",
+       "\n",
+       "import uuid\n",
        "from typing import Any\n",
        "\n",
        "from ragna.assistant import RagnaDemoAssistant\n",
@@ -305,7 +311,7 @@
        "\n",
        "    @classmethod\n",
        "    async def get_upload_info(\n",
-       "        cls, *, config: Config, user: str, id: str, name: str\n",
+       "        cls, *, config: Config, user: str, id: uuid.UUID, name: str\n",
        "    ) -> tuple[str, dict[str, Any], dict[str, Any]]:\n",
        "        if not PackageRequirement(\"boto3\").is_available():\n",
        "            raise RagnaException()\n",
@@ -373,9 +379,9 @@
    "id": "4088eb70-5f90-4458-bebe-a042b79a9502",
    "metadata": {},
    "source": [
-    "- `get_upload_info`: This method is called when the client hits the `/document/new` endpoint of the API. Here we generate the presigned URL and return the necessary information to the client so they can upload their file directly to S3\n",
-    "- `is_available`: This method is called when the client hits the `/chat/new` endpoint of the API. If the upload was not performed or failed, the API refuses to create a new chat with the specified document.\n",
-    "- `read`: This method is called when the client hits the `/chat/{id}/start` endpoint of the API, to store the content in the selected source storage\n",
+    "- `get_upload_info`: This method is called when the client hits the `/document` endpoint of the API. Here we generate the presigned URL and return the necessary information to the client so they can upload their file directly to S3\n",
+    "- `is_available`: This method is called when the client hits the `/chat` endpoint of the API. If the upload was not performed or failed, the API refuses to create a new chat with the specified document.\n",
+    "- `read`: This method is called when the client hits the `/chat/{id}/prepare` endpoint of the API, to store the content in the selected source storage\n",
     "\n",
     "From this point on, this notebook is a reduced version of the REST API example."
    ]
@@ -419,7 +425,25 @@
    "execution_count": 5,
    "id": "ae31bfb0-01f6-47c8-9733-ba1204b153d8",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:     Started server process [33447]\n",
+      "INFO:     Waiting for application startup.\n",
+      "INFO:     Application startup complete.\n",
+      "INFO:     Uvicorn running on http://127.0.0.1:31476 (Press CTRL+C to quit)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO:     127.0.0.1:51534 - \"GET / HTTP/1.1\" 200 OK\n"
+     ]
+    }
+   ],
    "source": [
     "import contextlib\n",
     "import subprocess\n",
@@ -435,7 +459,7 @@
     "start = time.time()\n",
     "while (time.time() - start) < timeout:\n",
     "    with contextlib.suppress(httpx.ConnectError):\n",
-    "        response = await client.get(f\"{URL}/health\")\n",
+    "        response = await client.get(URL)\n",
     "        if response.is_success:\n",
     "            break\n",
     "\n",
@@ -473,8 +497,9 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'data': {'token': 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VyIjoiUmFnbmEiLCJpZCI6ImIyZDRkNTgwLTA0OGQtNDBkMS1iMzljLTk1NGFmNDI5OTMwOSIsImV4cCI6MTY5Njg3Nzk4Mi4yOTA1NTk1fQ.0rVxqW0h5A3c8CJmp-roaRFUoJ7IAc4zUsQHJL2RuaU'},\n",
-      " 'document': {'id': 'b2d4d580-048d-40d1-b39c-954af4299309',\n",
+      "INFO:     127.0.0.1:51534 - \"GET /document?user=Ragna&name=document0.txt HTTP/1.1\" 200 OK\n",
+      "{'data': {'token': 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VyIjoiUmFnbmEiLCJpZCI6ImFkMWIwMWI1LTM0MjMtNDUxZS1hOWVlLTVlY2U5M2RlNGVlNiIsImV4cCI6MTY5NzExNTAxMC4yOTIwMTE1fQ.UOdIsvuWGV5NluQd1LrcMGjptLUnwmCSQIxn8UOEiUo'},\n",
+      " 'document': {'id': 'ad1b01b5-3423-451e-a9ee-5ece93de4ee6',\n",
       "              'name': 'document0.txt'},\n",
       " 'url': 'http://127.0.0.1:31476/document'}\n"
      ]
@@ -496,7 +521,15 @@
    "execution_count": 8,
    "id": "fdc5c6f4-bcf8-4902-b441-768b6e3a66bc",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO:     127.0.0.1:51534 - \"POST /document HTTP/1.1\" 200 OK\n"
+     ]
+    }
+   ],
    "source": [
     "response = await client.post(\n",
     "    document_info[\"url\"],\n",
@@ -512,12 +545,22 @@
    "id": "f0d769ff-d721-460a-8888-5b033eb4a909",
    "metadata": {},
    "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO:     127.0.0.1:51534 - \"GET /document?user=Ragna&name=document1.txt HTTP/1.1\" 200 OK\n",
+      "INFO:     127.0.0.1:51534 - \"POST /document HTTP/1.1\" 200 OK\n",
+      "INFO:     127.0.0.1:51534 - \"GET /document?user=Ragna&name=document2.txt HTTP/1.1\" 200 OK\n",
+      "INFO:     127.0.0.1:51534 - \"POST /document HTTP/1.1\" 200 OK\n"
+     ]
+    },
     {
      "data": {
       "text/plain": [
-       "[{'id': 'b2d4d580-048d-40d1-b39c-954af4299309', 'name': 'document0.txt'},\n",
-       " {'id': '525728e0-d478-4bce-8fdc-9f8b92f0cfe2', 'name': 'document1.txt'},\n",
-       " {'id': '308a261b-47a8-43b6-80aa-4c3c12549f99', 'name': 'document2.txt'}]"
+       "[{'id': 'ad1b01b5-3423-451e-a9ee-5ece93de4ee6', 'name': 'document0.txt'},\n",
+       " {'id': '7d901c1f-3638-4057-a1d2-6238e0a7601f', 'name': 'document1.txt'},\n",
+       " {'id': '19976a75-daf8-492c-8c9e-5075a9b7ca12', 'name': 'document2.txt'}]"
       ]
      },
      "execution_count": 9,
@@ -548,10 +591,17 @@
    "id": "ae364aec-e63b-4f00-8b55-217148d6df24",
    "metadata": {},
    "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO:     127.0.0.1:51534 - \"POST /chats?user=Ragna HTTP/1.1\" 200 OK\n"
+     ]
+    },
     {
      "data": {
       "text/plain": [
-       "'http://127.0.0.1:31476/chats/29e99963-08e9-4006-8b07-2f76e499a9c3'"
+       "'http://127.0.0.1:31476/chats/0a0fb603-77b8-44c8-8b18-9697be529af6'"
       ]
      },
      "execution_count": 10,
@@ -566,7 +616,7 @@
     "        params={\"user\": USER},\n",
     "        json={\n",
     "            \"name\": \"Ragna REST API example\",\n",
-    "            \"document_ids\": [d[\"id\"] for d in documents],\n",
+    "            \"documents\": documents,\n",
     "            \"source_storage\": \"Ragna/DemoSourceStorage\",\n",
     "            \"assistant\": \"Ragna/DemoAssistant\",\n",
     "            \"params\": {},\n",
@@ -588,6 +638,8 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
+      "INFO:     127.0.0.1:51534 - \"POST /chats/0a0fb603-77b8-44c8-8b18-9697be529af6/prepare?user=Ragna HTTP/1.1\" 200 OK\n",
+      "INFO:     127.0.0.1:51534 - \"POST /chats/0a0fb603-77b8-44c8-8b18-9697be529af6/answer?user=Ragna&prompt=Hello%20World%21 HTTP/1.1\" 200 OK\n",
       "I can't really help you with your prompt:\n",
       "\n",
       "> Hello World!\n",
@@ -601,7 +653,7 @@
     }
    ],
    "source": [
-    "await client.post(f\"{CHAT_URL}/start\", params={\"user\": USER})\n",
+    "await client.post(f\"{CHAT_URL}/prepare\", params={\"user\": USER})\n",
     "answer = (\n",
     "    await client.post(\n",
     "        f\"{CHAT_URL}/answer\", params={\"user\": USER, \"prompt\": \"Hello World!\"}\n",
diff --git a/ragna/__init__.py b/ragna/__init__.py
index 80e14ff7..d159d831 100644
--- a/ragna/__init__.py
+++ b/ragna/__init__.py
@@ -23,21 +23,21 @@ def demo_config():
 demo_config = demo_config()
 
 
-def builtin_config():
-    from ragna.core import Assistant, SourceStorage
-
-    builtin_config = Config()
-    builtin_config.state_database_url = (
-        f"sqlite:///{builtin_config.local_cache_root}/ragna.db"
-    )
-    builtin_config.queue_database_url = str(builtin_config.local_cache_root / "queue")
-
-    for module, cls in [(source_storage, SourceStorage), (assistant, Assistant)]:
-        for obj in module.__dict__.values():
-            if isinstance(obj, type) and issubclass(obj, cls) and obj.is_available():
-                builtin_config.register_component(obj)
-
-    return builtin_config
-
-
-builtin_config = builtin_config()
+# def builtin_config():
+#     from ragna.core import Assistant, SourceStorage
+#
+#     builtin_config = Config()
+#     builtin_config.state_database_url = (
+#         f"sqlite:///{builtin_config.local_cache_root}/ragna.db"
+#     )
+#     builtin_config.queue_database_url = str(builtin_config.local_cache_root / "queue")
+#
+#     for module, cls in [(source_storage, SourceStorage), (assistant, Assistant)]:
+#         for obj in module.__dict__.values():
+#             if isinstance(obj, type) and issubclass(obj, cls) and obj.is_available():
+#                 builtin_config.register_component(obj)
+#
+#     return builtin_config
+#
+#
+# builtin_config = builtin_config()
diff --git a/ragna/_api/core.py b/ragna/_api/core.py
index 18ba376c..4446ef8f 100644
--- a/ragna/_api/core.py
+++ b/ragna/_api/core.py
@@ -1,15 +1,16 @@
 import functools
+import uuid
 from typing import Annotated
-from uuid import UUID
 
 import aiofiles
 from fastapi import Depends, FastAPI, Form, HTTPException, UploadFile
 
 import ragna
+import ragna.core
 
-from ragna.core import Chat, LocalDocument, RagnaException, RagnaId
+from ragna.core import Rag, RagnaException
 
-from . import schemas
+from . import database, schemas
 
 
 def process_ragna_exception(afn):
@@ -33,10 +34,12 @@ async def wrapper(*args, **kwargs):
     return wrapper
 
 
-def api(rag):
+def api(config):
+    rag = Rag(config)
+
     app = FastAPI()
 
-    @app.get("/health")
+    @app.get("/")
     @process_ragna_exception
     async def health() -> str:
         return ragna.__version__
@@ -55,96 +58,159 @@ async def get_components(_: UserDependency) -> schemas.Components:
             assistants=list(rag.config.registered_assistant_classes),
         )
 
+    make_session = database.get_sessionmaker(config.state_database_url)
+
+    def get_session():
+        session = make_session()
+        try:
+            yield session
+        finally:
+            session.close()
+
+    SessionDependency = Annotated[database.Session, Depends(get_session)]
+
     @app.get("/document")
     @process_ragna_exception
     async def get_document_upload_info(
+        session: SessionDependency,
         user: UserDependency,
         name: str,
     ) -> schemas.DocumentUploadInfo:
-        id = RagnaId.make()
-        url, data, metadata = await rag.config.document_class.get_upload_info(
-            config=rag.config, user=user, id=id, name=name
-        )
-        rag._add_document(user=user, id=id, name=name, metadata=metadata)
-        return schemas.DocumentUploadInfo(
-            url=url, data=data, document=schemas.Document(id=id, name=name)
+        document = schemas.Document(name=name)
+        url, data, metadata = await config.document_class.get_upload_info(
+            config=config, user=user, id=document.id, name=document.name
         )
+        database.add_document(session, user=user, document=document, metadata=metadata)
+        return schemas.DocumentUploadInfo(url=url, data=data, document=document)
 
     @app.post("/document")
     @process_ragna_exception
     async def upload_document(
-        token: Annotated[str, Form()], file: UploadFile
+        session: SessionDependency, token: Annotated[str, Form()], file: UploadFile
     ) -> schemas.Document:
-        if not issubclass(rag.config.document_class, LocalDocument):
+        if not issubclass(rag.config.document_class, ragna.core.LocalDocument):
             raise HTTPException(
                 status_code=400,
                 detail="Ragna configuration does not support local upload",
             )
 
-        user, id = rag.config.document_class._decode_upload_token(
+        user, id = ragna.core.LocalDocument._decode_upload_token(
             token, secret=rag.config.upload_token_secret
         )
-        document = rag._get_document(user=user, id=id)
+        document, metadata = database.get_document(session, user=user, id=id)
 
-        document.path.parent.mkdir(parents=True, exist_ok=True)
-        async with aiofiles.open(document.path, "wb") as document_file:
+        core_document = ragna.core.LocalDocument(
+            id=document.id, name=document.name, metadata=metadata
+        )
+        core_document.path.parent.mkdir(parents=True, exist_ok=True)
+        async with aiofiles.open(core_document.path, "wb") as document_file:
             while content := await file.read(1024):
                 await document_file.write(content)
 
-        return schemas.Document(id=id, name=document.name)
+        return document
+
+    def schema_to_core_chat(
+        session, *, user: str, chat: schemas.Chat
+    ) -> ragna.core.Chat:
+        core_chat = rag.chat(
+            documents=[
+                rag.config.document_class(
+                    id=document.id,
+                    name=document.name,
+                    metadata=database.get_document(
+                        session,
+                        user=user,
+                        id=document.id,
+                    )[1],
+                )
+                for document in chat.metadata.documents
+            ],
+            source_storage=chat.metadata.source_storage,
+            assistant=chat.metadata.assistant,
+            user=user,
+            chat_id=chat.id,
+            chat_name=chat.metadata.name,
+            **chat.metadata.params,
+        )
+        # FIXME: We need to reconstruct the previous messages here. Right now this is
+        #  not needed, because the chat itself never accesses past messages. However,
+        #  if we implement a chat history feature, i.e. passing past messages to
+        #  the assistant, this becomes crucial.
+        core_chat.messages = []
+        core_chat._prepared = chat.prepared
+
+        return core_chat
 
     @app.post("/chats")
     @process_ragna_exception
     async def create_chat(
-        *, user: UserDependency, chat_metadata: schemas.ChatMetadataCreate
+        session: SessionDependency,
+        user: UserDependency,
+        chat_metadata: schemas.ChatMetadata,
     ) -> schemas.Chat:
-        return schemas.Chat.from_core_chat(
-            await rag.new_chat(
-                user=user,
-                name=chat_metadata.name,
-                documents=chat_metadata.document_ids,
-                source_storage=chat_metadata.source_storage,
-                assistant=chat_metadata.assistant,
-                **chat_metadata.params,
-            )
-        )
-
-    @app.get("/chats")
-    @process_ragna_exception
-    async def get_chats(user: UserDependency) -> list[schemas.Chat]:
-        return [schemas.Chat.from_core_chat(chat) for chat in rag._get_chats(user=user)]
-
-    async def _get_id(id: UUID) -> RagnaId:
-        return RagnaId.from_uuid(id)
+        chat = schemas.Chat(metadata=chat_metadata)
 
-    IdDependency = Annotated[RagnaId, Depends(_get_id)]
+        # Although we don't need the actual ragna.core.Chat object here,
+        # we use it to validate the documents and metadata.
+        schema_to_core_chat(session, user=user, chat=chat)
 
-    async def _get_chat(*, user: UserDependency, id: IdDependency) -> Chat:
-        return rag._get_chat(user=user, id=id)
+        database.add_chat(session, user=user, chat=chat)
+        return chat
 
-    ChatDependency = Annotated[Chat, Depends(_get_chat, use_cache=False)]
+    @app.get("/chats")
+    @process_ragna_exception
+    async def get_chats(
+        session: SessionDependency, user: UserDependency
+    ) -> list[schemas.Chat]:
+        return database.get_chats(session, user=user)
 
     @app.get("/chats/{id}")
     @process_ragna_exception
-    async def get_chat(chat: ChatDependency) -> schemas.Chat:
-        return schemas.Chat.from_core_chat(chat)
+    async def get_chat(
+        session: SessionDependency, user: UserDependency, id: uuid.UUID
+    ) -> schemas.Chat:
+        return database.get_chat(session, user=user, id=id)
 
-    @app.post("/chats/{id}/start")
+    @app.post("/chats/{id}/prepare")
     @process_ragna_exception
-    async def start_chat(chat: ChatDependency) -> schemas.Chat:
-        return schemas.Chat.from_core_chat(await chat.start())
+    async def prepare_chat(
+        session: SessionDependency, user: UserDependency, id: uuid.UUID
+    ) -> schemas.MessageOutput:
+        chat = database.get_chat(session, user=user, id=id)
 
-    @app.post("/chats/{id}/close")
-    @process_ragna_exception
-    async def close_chat(chat: ChatDependency) -> schemas.Chat:
-        return schemas.Chat.from_core_chat(await chat.close())
+        core_chat = schema_to_core_chat(session, user=user, chat=chat)
+        welcome = schemas.Message.from_core(await core_chat.prepare())
+        chat.prepared = True
+        chat.messages.append(welcome)
+
+        database.update_chat(session, user=user, chat=chat)
+
+        return schemas.MessageOutput(message=welcome, chat=chat)
 
     @app.post("/chats/{id}/answer")
     @process_ragna_exception
-    async def answer(chat: ChatDependency, prompt: str) -> schemas.AnswerOutput:
-        return schemas.AnswerOutput(
-            message=schemas.Message.from_core_message(await chat.answer(prompt)),
-            chat=schemas.Chat.from_core_chat(chat),
+    async def answer(
+        session: SessionDependency, user: UserDependency, id: uuid.UUID, prompt: str
+    ) -> schemas.MessageOutput:
+        chat = database.get_chat(session, user=user, id=id)
+        chat.messages.append(
+            schemas.Message(content=prompt, role=ragna.core.MessageRole.USER)
         )
 
+        core_chat = schema_to_core_chat(session, user=user, chat=chat)
+
+        answer = schemas.Message.from_core(await core_chat.answer(prompt))
+        chat.messages.append(answer)
+
+        database.update_chat(session, user=user, chat=chat)
+
+        return schemas.MessageOutput(message=answer, chat=chat)
+
+    @app.delete("/chats/{id}")
+    @process_ragna_exception
+    async def delete_chat(
+        session: SessionDependency, user: UserDependency, id: uuid.UUID
+    ) -> None:
+        database.delete_chat(session, user=user, id=id)
+
     return app
diff --git a/ragna/_api/database.py b/ragna/_api/database.py
new file mode 100644
index 00000000..73507b79
--- /dev/null
+++ b/ragna/_api/database.py
@@ -0,0 +1,215 @@
+from __future__ import annotations
+
+import functools
+
+import uuid
+from typing import Any, Callable
+
+from sqlalchemy import create_engine, select
+from sqlalchemy.orm import Session, sessionmaker as _sessionmaker
+
+from ragna.core import RagnaException
+
+from . import orm, schemas
+
+
+def get_sessionmaker(database_url: str) -> Callable[[], Session]:
+    engine = create_engine(database_url, connect_args=dict(check_same_thread=False))
+    orm.Base.metadata.create_all(bind=engine)
+    return _sessionmaker(autocommit=False, autoflush=False, bind=engine)
+
+
+@functools.lru_cache(maxsize=1024)
+def _get_user_id(session: Session, username: str) -> uuid.UUID:
+    user = session.execute(
+        select(orm.User).where(orm.User.name == username)
+    ).scalar_one_or_none()
+
+    if user is None:
+        # Add a new user if the current username is not registered yet. Since this is
+        # behind the authentication layer, we don't need any extra security here.
+        user = orm.User(id=uuid.uuid4(), name=username)
+        session.add(user)
+        session.commit()
+
+    return user.id
+
+
+def add_document(
+    session: Session, *, user: str, document: schemas.Document, metadata: dict[str, Any]
+) -> None:
+    session.add(
+        orm.Document(
+            id=document.id,
+            user_id=_get_user_id(session, user),
+            name=document.name,
+            metadata_=metadata,
+        )
+    )
+    session.commit()
+
+
+def _orm_to_schema_document(document: orm.Document) -> schemas.Document:
+    return schemas.Document(id=document.id, name=document.name)
+
+
+@functools.lru_cache(maxsize=1024)
+def get_document(
+    session: Session, *, user: str, id: uuid.UUID
+) -> tuple[schemas.Document, dict[str, Any]]:
+    document = session.execute(
+        select(orm.Document).where(
+            (orm.Document.user_id == _get_user_id(session, user))
+            & (orm.Document.id == id)
+        )
+    ).scalar_one_or_none()
+    return _orm_to_schema_document(document), document.metadata_
+
+
+def add_chat(session: Session, *, user: str, chat: schemas.Chat):
+    document_ids = {document.id for document in chat.metadata.documents}
+    documents = (
+        session.execute(select(orm.Document).where(orm.Document.id.in_(document_ids)))
+        .scalars()
+        .all()
+    )
+    if len(documents) != len(document_ids):
+        raise RagnaException(
+            set(document_ids) - {document.id for document in documents}
+        )
+    session.add(
+        orm.Chat(
+            id=chat.id,
+            user_id=_get_user_id(session, user),
+            name=chat.metadata.name,
+            documents=documents,
+            source_storage=chat.metadata.source_storage,
+            assistant=chat.metadata.assistant,
+            params=chat.metadata.params,
+            prepared=chat.prepared,
+        )
+    )
+    session.commit()
+
+
+def _orm_to_schema_chat(chat: orm.Chat) -> schemas.Chat:
+    documents = [
+        schemas.Document(id=document.id, name=document.name)
+        for document in chat.documents
+    ]
+    messages = [
+        schemas.Message(
+            id=message.id,
+            role=message.role,
+            content=message.content,
+            sources=[
+                schemas.Source(
+                    id=source.id,
+                    document=_orm_to_schema_document(source.document),
+                    location=source.location,
+                )
+                for source in message.sources
+            ],
+            timestamp=message.timestamp,
+        )
+        for message in chat.messages
+    ]
+    return schemas.Chat(
+        id=chat.id,
+        metadata=schemas.ChatMetadata(
+            name=chat.name,
+            documents=documents,
+            source_storage=chat.source_storage,
+            assistant=chat.assistant,
+            params=chat.params,
+        ),
+        messages=messages,
+        prepared=chat.prepared,
+    )
+
+
+def get_chats(session: Session, *, user: str) -> list[schemas.Chat]:
+    return [
+        _orm_to_schema_chat(chat)
+        for chat in session.execute(
+            select(orm.Chat).where(orm.Chat.user_id == _get_user_id(session, user))
+        )
+        .scalars()
+        .all()
+    ]
+
+
+def _get_orm_chat(session: Session, *, user: str, id: uuid.UUID) -> orm.Chat:
+    chat = session.execute(
+        select(orm.Chat).where(
+            (orm.Chat.id == id) & (orm.Chat.user_id == _get_user_id(session, user))
+        )
+    ).scalar_one_or_none()
+    if chat is None:
+        raise RagnaException()
+    return chat
+
+
+def get_chat(session: Session, *, user: str, id: uuid.UUID) -> schemas.Chat:
+    return _orm_to_schema_chat(_get_orm_chat(session, user=user, id=id))
+
+
+def _schema_to_orm_source(session: Session, source: schemas.Source) -> orm.Source:
+    orm_source = session.execute(
+        select(orm.Source).where(orm.Source.id == source.id)
+    ).scalar_one_or_none()
+
+    if orm_source is None:
+        orm_source = orm.Source(
+            id=source.id,
+            document_id=source.document.id,
+            location=source.location,
+        )
+        session.add(orm_source)
+        session.commit()
+        session.refresh(orm_source)
+
+    return orm_source
+
+
+def _schema_to_orm_message(
+    session: Session, chat_id: uuid.UUID, message: schemas.Message
+) -> orm.Message:
+    orm_message = session.execute(
+        select(orm.Message).where(orm.Message.id == message.id)
+    ).scalar_one_or_none()
+    if orm_message is None:
+        orm_message = orm.Message(
+            id=message.id,
+            chat_id=chat_id,
+            content=message.content,
+            role=message.role,
+            sources=[
+                _schema_to_orm_source(session, source=source)
+                for source in message.sources
+            ],
+            timestamp=message.timestamp,
+        )
+        session.add(orm_message)
+        session.commit()
+        session.refresh(orm_message)
+
+    return orm_message
+
+
+def update_chat(session: Session, user: str, chat: schemas.Chat) -> None:
+    orm_chat = _get_orm_chat(session, user=user, id=chat.id)
+
+    orm_chat.prepared = chat.prepared
+    orm_chat.messages = [
+        _schema_to_orm_message(session, chat_id=chat.id, message=message)
+        for message in chat.messages
+    ]
+
+    session.commit()
+
+
+def delete_chat(session: Session, user: str, id: uuid.UUID) -> None:
+    orm_chat = _get_orm_chat(session, user=user, id=id)
+    session.delete(orm_chat)
+    session.commit()
diff --git a/ragna/_api/orm.py b/ragna/_api/orm.py
new file mode 100644
index 00000000..41f9244d
--- /dev/null
+++ b/ragna/_api/orm.py
@@ -0,0 +1,107 @@
+from sqlalchemy import Column, ForeignKey, Table, types
+from sqlalchemy.orm import DeclarativeBase, relationship
+
+from ragna.core import MessageRole
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+# FIXME: Do we actually need this table? If we are sure that a user is unique and has to
+#  be authenticated from the API layer, it seems having an extra mapping here is not
+#  needed?
+class User(Base):
+    __tablename__ = "users"
+
+    id = Column(types.Uuid, primary_key=True)
+    name = Column(types.String)
+
+
+document_chat_association_table = Table(
+    "document_chat_association_table",
+    Base.metadata,
+    Column("document_id", ForeignKey("documents.id"), primary_key=True),
+    Column("chat_id", ForeignKey("chats.id"), primary_key=True),
+)
+
+
+class Document(Base):
+    __tablename__ = "documents"
+
+    id = Column(types.Uuid, primary_key=True)
+    user_id = Column(ForeignKey("users.id"))
+    name = Column(types.String)
+    # Mind the trailing underscore here. Unfortunately, this is necessary, because
+    # metadata without the underscore is reserved by SQLAlchemy
+    metadata_ = Column(types.JSON)
+    chats = relationship(
+        "Chat",
+        secondary=document_chat_association_table,
+        back_populates="documents",
+    )
+    sources = relationship(
+        "Source",
+        back_populates="document",
+    )
+
+
+class Chat(Base):
+    __tablename__ = "chats"
+
+    id = Column(types.Uuid, primary_key=True)
+    user_id = Column(ForeignKey("users.id"))
+    name = Column(types.String)
+    documents = relationship(
+        "Document",
+        secondary=document_chat_association_table,
+        back_populates="chats",
+    )
+    source_storage = Column(types.String)
+    assistant = Column(types.String)
+    params = Column(types.JSON)
+    messages = relationship("Message", cascade="all, delete")
+    prepared = Column(types.Boolean)
+
+
+source_message_association_table = Table(
+    "source_message_association_table",
+    Base.metadata,
+    Column("source_id", ForeignKey("sources.id"), primary_key=True),
+    Column("message_id", ForeignKey("messages.id"), primary_key=True),
+)
+
+
+class Source(Base):
+    __tablename__ = "sources"
+
+    # This is not a UUID column as all of the other IDs, because we don't control its
+    # generation. It is generated as part of ragna.core.SourceStorage.retrieve and using
+    # a string here doesn't impose unnecessary constraints on the user.
+    id = Column(types.String, primary_key=True)
+
+    document_id = Column(ForeignKey("documents.id"))
+    document = relationship("Document", back_populates="sources")
+
+    location = Column(types.String)
+
+    messages = relationship(
+        "Message",
+        secondary=source_message_association_table,
+        back_populates="sources",
+    )
+
+
+class Message(Base):
+    __tablename__ = "messages"
+
+    id = Column(types.Uuid, primary_key=True)
+    chat_id = Column(ForeignKey("chats.id"))
+    content = Column(types.String)
+    role = Column(types.Enum(MessageRole))
+    sources = relationship(
+        "Source",
+        secondary=source_message_association_table,
+        back_populates="messages",
+    )
+    timestamp = Column(types.DateTime)
diff --git a/ragna/_api/schemas.py b/ragna/_api/schemas.py
index d00b0fcd..031e9be8 100644
--- a/ragna/_api/schemas.py
+++ b/ragna/_api/schemas.py
@@ -1,21 +1,24 @@
 from __future__ import annotations
 
 import datetime
-from uuid import UUID
+import uuid
 
-from pydantic import BaseModel, HttpUrl, validator
-
-import ragna
+from pydantic import BaseModel, Field, HttpUrl
 
 import ragna.core
 
 
+class Components(BaseModel):
+    source_storages: list[str]
+    assistants: list[str]
+
+
 class Document(BaseModel):
-    id: ragna.core.RagnaId
+    id: uuid.UUID = Field(default_factory=uuid.uuid4)
     name: str
 
     @classmethod
-    def from_core_document(cls, document: ragna.core.Document) -> Document:
+    def from_core(cls, document: ragna.core.Document) -> Document:
         return cls(
             id=document.id,
             name=document.name,
@@ -29,91 +32,53 @@ class DocumentUploadInfo(BaseModel):
 
 
 class Source(BaseModel):
-    id: ragna.core.RagnaId
+    # See orm.Source on why this is not a UUID
+    id: str
     document: Document
     location: str
 
     @classmethod
-    def from_core_source(cls, source: ragna.core.Source) -> Source:
+    def from_core(cls, source: ragna.core.Source) -> Source:
         return cls(
             id=source.id,
-            document=Document(id=source.document_id, name=source.document_name),
+            document=Document.from_core(source.document),
             location=source.location,
         )
 
 
 class Message(BaseModel):
-    id: ragna.core.RagnaId
-    role: ragna.core.MessageRole
+    id: uuid.UUID = Field(default_factory=uuid.uuid4)
     content: str
-    sources: list[Source]
-    timestamp: datetime.datetime
+    role: ragna.core.MessageRole
+    sources: list[Source] = Field(default_factory=list)
+    timestamp: datetime.datetime = Field(
+        default_factory=lambda: datetime.datetime.now(tz=datetime.timezone.utc)
+    )
 
     @classmethod
-    def from_core_message(cls, message: ragna.core.Message) -> Message:
+    def from_core(cls, message: ragna.core.Message) -> Message:
         return cls(
-            id=message.id,
-            role=message.role,
             content=message.content,
-            sources=[Source.from_core_source(s) for s in message.sources],
-            timestamp=message.timestamp,
+            role=message.role,
+            sources=[Source.from_core(source) for source in message.sources],
         )
 
 
-class ChatMetadataBase(BaseModel):
+class ChatMetadata(BaseModel):
     name: str
     source_storage: str
     assistant: str
     params: dict
-
-
-class ChatMetadataCreate(ChatMetadataBase):
-    # For some reason list[RagnaId] does not work and will get parsed into list[UUID].
-    # Thus, we use a validator below to do the conversion.
-    document_ids: list[UUID]
-
-    @validator("document_ids")
-    def uuid_to_ragna_id(cls, document_ids: list[UUID]) -> list[ragna.core.RagnaId]:
-        return [ragna.core.RagnaId.from_uuid(u) for u in document_ids]
-
-
-class ChatMetadata(ChatMetadataBase):
     documents: list[Document]
 
-    @classmethod
-    def from_core_chat(cls, chat: ragna.core.Chat) -> ChatMetadata:
-        return cls(
-            name=chat.name,
-            documents=[Document.from_core_document(d) for d in chat.documents],
-            source_storage=chat.source_storage.display_name(),
-            assistant=chat.assistant.display_name(),
-            params=chat.params,
-        )
-
 
 class Chat(BaseModel):
-    id: ragna.core.RagnaId
+    id: uuid.UUID = Field(default_factory=uuid.uuid4)
     metadata: ChatMetadata
-    messages: list[Message]
-    started: bool
-    closed: bool
-
-    @classmethod
-    def from_core_chat(cls, chat: ragna.core.Chat) -> Chat:
-        return cls(
-            id=chat.id,
-            metadata=ChatMetadata.from_core_chat(chat),
-            messages=[Message.from_core_message(m) for m in chat.messages],
-            started=chat._started,
-            closed=chat._closed,
-        )
+    messages: list[Message] = Field(default_factory=list)
+    prepared: bool = False
 
 
-class AnswerOutput(BaseModel):
+class MessageOutput(BaseModel):
     message: Message
     chat: Chat
-
-
-class Components(BaseModel):
-    source_storages: list[str]
-    assistants: list[str]
diff --git a/ragna/_cli.py b/ragna/_cli.py
index dfb9791a..963d6545 100644
--- a/ragna/_cli.py
+++ b/ragna/_cli.py
@@ -9,7 +9,7 @@
 
 import ragna
 
-from ragna.core import Config, EnvVarRequirement, PackageRequirement, Rag, Requirement
+from ragna.core import Config, EnvVarRequirement, PackageRequirement, Requirement
 from ragna.core._queue import Queue
 
 app = typer.Typer(
@@ -98,7 +98,7 @@ def _yes_or_no(condition):
 
 
 @app.command(help="Start Ragna API")
-def api(*, config: ConfigAnnotated = "ragna.builtin_config"):
+def api(*, config: ConfigAnnotated = "ragna.demo_config"):
     required_packages = [
         package
         for package in ["fastapi", "uvicorn"]
@@ -112,16 +112,14 @@ def api(*, config: ConfigAnnotated = "ragna.builtin_config"):
 
     from ragna._api import api
 
-    rag = Rag(config=config)
-
     components = urlsplit(config.ragna_api_url)
-    uvicorn.run(api(rag), host=components.hostname, port=components.port)
+    uvicorn.run(api(config), host=components.hostname, port=components.port)
 
 
 @app.command(help="Start Ragna worker(s)")
 def worker(
     *,
-    config: ConfigAnnotated = "ragna.builtin_config",
+    config: ConfigAnnotated = "ragna.demo_config",
     num_workers: Annotated[int, typer.Option("--num-workers", "-n")] = 1,
 ):
     queue = Queue(config, load_components=True)
diff --git a/ragna/assistant/_api.py b/ragna/assistant/_api.py
index 03f57589..d7e80fc0 100644
--- a/ragna/assistant/_api.py
+++ b/ragna/assistant/_api.py
@@ -35,7 +35,8 @@ def __init__(self, config, *, num_retries: int = 2, retry_delay: float = 1.0):
         import httpx
 
         self._client = httpx.Client(
-            headers={"User-Agent": f"{ragna.__version__}/{self}"}
+            headers={"User-Agent": f"{ragna.__version__}/{self}"},
+            timeout=10,
         )
         self._num_retries = num_retries
         self._retry_delay = retry_delay
diff --git a/ragna/assistant/_demo.py b/ragna/assistant/_demo.py
index 3d2d76fe..37b503f5 100644
--- a/ragna/assistant/_demo.py
+++ b/ragna/assistant/_demo.py
@@ -34,7 +34,7 @@ def _markdown_answer(self):
     def _default_answer(self, prompt: str, sources: list[Source]) -> str:
         sources_display = []
         for source in sources:
-            source_display = f"- {source.document_name}"
+            source_display = f"- {source.document.name}"
             if source.location:
                 source_display += f", {source.location}"
             source_display += f": {textwrap.shorten(source.content, width=100)}"
diff --git a/ragna/core/__init__.py b/ragna/core/__init__.py
index b1dc3725..1d231427 100644
--- a/ragna/core/__init__.py
+++ b/ragna/core/__init__.py
@@ -1,4 +1,4 @@
-from ._core import RagnaException, RagnaId  # usort: skip
+from ._core import RagnaException  # usort: skip
 
 from ._config import Config
 from ._queue import task_config
@@ -9,7 +9,7 @@
 )  # usort: skip
 
 from ._document import Document, LocalDocument, Page, PageExtractor
-from ._source_storage import ReconstructedSource, Source, SourceStorage, Tokenizer
+from ._source_storage import Source, SourceStorage, Tokenizer
 from ._assistant import Assistant, Message, MessageRole  # usort: skip
 
 from ._rag import Chat, Rag  # usort: skip
diff --git a/ragna/core/_assistant.py b/ragna/core/_assistant.py
index 39e362db..79c609e7 100644
--- a/ragna/core/_assistant.py
+++ b/ragna/core/_assistant.py
@@ -1,11 +1,10 @@
 import abc
-import datetime
+
 import enum
 
-from typing import Optional
+from pydantic import BaseModel, Field
 
 from ._component import RagComponent
-from ._core import RagnaId
 from ._source_storage import Source
 
 
@@ -15,22 +14,12 @@ class MessageRole(enum.Enum):
     ASSISTANT = "assistant"
 
 
-class Message:
-    def __init__(
-        self,
-        *,
-        id: RagnaId,
-        content: str,
-        role: MessageRole,
-        sources: Optional[list[Source]] = None,
-    ):
-        self.id = id
-        self.content = content
-        self.role = role
-        self.sources = sources or []
-        self.timestamp = datetime.datetime.utcnow()
-
-    def __repr__(self):
+class Message(BaseModel):
+    content: str
+    role: MessageRole
+    sources: list[Source] = Field(default_factory=list)
+
+    def __str__(self):
         return self.content
 
 
diff --git a/ragna/core/_core.py b/ragna/core/_core.py
index 66949bc0..67ce7317 100644
--- a/ragna/core/_core.py
+++ b/ragna/core/_core.py
@@ -1,8 +1,3 @@
-from __future__ import annotations
-
-import uuid as uuid_
-
-
 class RagnaException(Exception):
     # The values below are sentinels to be used with the http_detail field.
     # This tells the API to use the event as detail
@@ -19,17 +14,3 @@ def __init__(self, event="", http_status_code=500, http_detail=None, **extra):
 
     def __str__(self):
         return ", ".join([self.event, *[f"{k}={v}" for k, v in self.extra.items()]])
-
-
-class RagnaId(uuid_.UUID):
-    @classmethod
-    def from_uuid(cls, uuid: uuid_.UUID) -> RagnaId:
-        return cls(int=uuid.int)
-
-    @classmethod
-    def is_valid_str(cls):
-        pass
-
-    @staticmethod
-    def make():
-        return RagnaId.from_uuid(uuid_.uuid4())
diff --git a/ragna/core/_document.py b/ragna/core/_document.py
index 6b376ebf..027a39cc 100644
--- a/ragna/core/_document.py
+++ b/ragna/core/_document.py
@@ -3,10 +3,11 @@
 import abc
 import dataclasses
 import time
+import uuid
 from pathlib import Path
 from typing import Any, Iterator, Optional, TYPE_CHECKING
 
-from ._core import RagnaException, RagnaId
+from ._core import RagnaException
 from ._requirement import PackageRequirement, Requirement, RequirementMixin
 
 if TYPE_CHECKING:
@@ -17,12 +18,12 @@ class Document(abc.ABC):
     def __init__(
         self,
         *,
-        id: Optional[RagnaId] = None,
+        id: Optional[uuid.UUID] = None,
         name: str,
         metadata: dict[str, Any],
         page_extractor: Optional[PageExtractor] = None,
     ):
-        self.id = id
+        self.id = id or uuid.uuid4()
         self.name = name
         self.metadata = metadata
 
@@ -37,7 +38,7 @@ def __init__(
     @classmethod
     @abc.abstractmethod
     async def get_upload_info(
-        cls, *, config: Config, user: str, id: str, name: str
+        cls, *, config: Config, user: str, id: uuid.UUID, name: str
     ) -> tuple[str, dict[str, Any], dict[str, Any]]:
         pass
 
@@ -84,7 +85,7 @@ def __init__(
 
     @classmethod
     async def get_upload_info(
-        cls, *, config: Config, user: str, id: RagnaId, name: str
+        cls, *, config: Config, user: str, id: uuid.UUID, name: str
     ) -> tuple[str, dict[str, Any], dict[str, Any]]:
         if not PackageRequirement("PyJWT").is_available():
             raise RagnaException()
@@ -107,7 +108,7 @@ async def get_upload_info(
         return url, data, metadata
 
     @classmethod
-    def _decode_upload_token(cls, token: str, *, secret: str) -> tuple[str, RagnaId]:
+    def _decode_upload_token(cls, token: str, *, secret: str) -> tuple[str, uuid.UUID]:
         import jwt
 
         try:
@@ -120,7 +121,7 @@ def _decode_upload_token(cls, token: str, *, secret: str) -> tuple[str, RagnaId]
             raise RagnaException(
                 "Token expired", http_status_code=401, http_detail=RagnaException.EVENT
             )
-        return payload["user"], RagnaId(payload["id"])
+        return payload["user"], uuid.UUID(payload["id"])
 
     @property
     def path(self) -> Path:
diff --git a/ragna/core/_orm.py b/ragna/core/_orm.py
deleted file mode 100644
index e1ff9149..00000000
--- a/ragna/core/_orm.py
+++ /dev/null
@@ -1,120 +0,0 @@
-from sqlalchemy import Column, ForeignKey, Table, types
-from sqlalchemy.orm import DeclarativeBase, relationship
-
-from ragna.core import MessageRole, RagnaId
-
-
-class Id(types.TypeDecorator):
-    impl = types.Uuid
-
-    cache_ok = True
-
-    def process_bind_param(self, value, dialect):
-        return value
-
-    def process_result_value(self, value, dialect):
-        if value is None:
-            return None
-        return RagnaId.from_uuid(value)
-
-
-class Base(DeclarativeBase):
-    pass
-
-
-# FIXME: Do we actually need this table? If we are sure that a user is unique and has to
-#  be authenticated from the API layer, it seems having an extra mapping here is not
-#  needed?
-class UserState(Base):
-    __tablename__ = "users"
-
-    id = Column(Id, primary_key=True)
-    name = Column(types.String)
-
-
-document_chat_state_association_table = Table(
-    "document_chat_state_association_table",
-    Base.metadata,
-    Column("document_id", ForeignKey("documents.id"), primary_key=True),
-    Column("chat_id", ForeignKey("chats.id"), primary_key=True),
-)
-
-
-class DocumentState(Base):
-    __tablename__ = "documents"
-
-    id = Column(Id, primary_key=True)
-    user_id = Column(ForeignKey("users.id"))
-    name = Column(types.String)
-    # Mind the trailing underscore here. Unfortunately, this is necessary, because
-    # metadata without the underscore is reserved by SQLAlchemy
-    metadata_ = Column(types.JSON)
-    chat_states = relationship(
-        "ChatState",
-        secondary=document_chat_state_association_table,
-        back_populates="document_states",
-    )
-    source_states = relationship(
-        "SourceState",
-        back_populates="document_state",
-    )
-
-
-class ChatState(Base):
-    __tablename__ = "chats"
-
-    id = Column(Id, primary_key=True)
-    user_id = Column(ForeignKey("users.id"))
-    name = Column(types.String)
-    document_states = relationship(
-        "DocumentState",
-        secondary=document_chat_state_association_table,
-        back_populates="chat_states",
-    )
-    source_storage = Column(types.String)
-    assistant = Column(types.String)
-    params = Column(types.JSON)
-    message_states = relationship("MessageState")
-    started = Column(types.Boolean)
-    closed = Column(types.Boolean)
-
-
-source_message_state_association_table = Table(
-    "source_message_state_association_table",
-    Base.metadata,
-    Column("source_id", ForeignKey("sources.id"), primary_key=True),
-    Column("message_id", ForeignKey("messages.id"), primary_key=True),
-)
-
-
-class SourceState(Base):
-    __tablename__ = "sources"
-
-    id = Column(Id, primary_key=True)
-
-    document_id = Column(ForeignKey("documents.id"))
-    document_state = relationship("DocumentState", back_populates="source_states")
-
-    location = Column(types.String)
-
-    message_states = relationship(
-        "MessageState",
-        secondary=source_message_state_association_table,
-        back_populates="source_states",
-    )
-
-
-class MessageState(Base):
-    __tablename__ = "messages"
-
-    id = Column(Id, primary_key=True)
-    chat_id = Column(ForeignKey("chats.id"))
-    content = Column(types.String)
-    role = Column(types.Enum(MessageRole))
-    source_id = Column(ForeignKey("sources.id"))
-    source_states = relationship(
-        "SourceState",
-        secondary=source_message_state_association_table,
-        back_populates="message_states",
-    )
-    timestamp = Column(types.DateTime)
diff --git a/ragna/core/_queue.py b/ragna/core/_queue.py
index 74e07c96..2f5dee9a 100644
--- a/ragna/core/_queue.py
+++ b/ragna/core/_queue.py
@@ -1,5 +1,5 @@
 import itertools
-from typing import Optional, Type, Union
+from typing import Optional, Type, TypeVar, Union
 from urllib.parse import urlsplit
 
 import huey.api
@@ -8,7 +8,6 @@
 
 from ._component import RagComponent
 from ._config import Config
-
 from ._core import RagnaException
 from ._requirement import PackageRequirement
 
@@ -35,13 +34,14 @@ def execute(self):
         return execute(*self.args)
 
 
+T = TypeVar("T", bound=RagComponent)
+
+
 class Queue:
     def __init__(self, config: Config, *, load_components: Optional[bool]):
         self._config = config
         self._huey = self._load_huey(config.queue_database_url)
 
-        if load_components is None:
-            load_components = isinstance(self._huey, huey.MemoryHuey)
         for component in itertools.chain(
             config.registered_source_storage_classes.values(),
             config.registered_assistant_classes.values(),
@@ -82,10 +82,13 @@ def _load_huey(self, url: Optional[str]):
 
     def parse_component(
         self,
-        component: Union[Type[RagComponent], RagComponent, str],
+        component: Union[Type[T], T, str],
         *,
-        load: bool = False,
-    ) -> Type[RagComponent]:
+        load: Optional[bool] = None,
+    ) -> Type[T]:
+        if load is None:
+            load = isinstance(self._huey, huey.MemoryHuey)
+
         if isinstance(component, type) and issubclass(component, RagComponent):
             cls = component
             instance = None
diff --git a/ragna/core/_rag.py b/ragna/core/_rag.py
index 8c509124..bd87b49b 100644
--- a/ragna/core/_rag.py
+++ b/ragna/core/_rag.py
@@ -1,23 +1,22 @@
 from __future__ import annotations
 
 import datetime
+
 import itertools
+import os
+import uuid
 from collections import defaultdict
-from typing import Any, Optional, Sequence, Type, TypeVar, Union
+from typing import Any, Iterable, Optional, Type, Union
 
-from pydantic import BaseModel, create_model, Extra
+from pydantic import BaseModel, create_model, Extra, Field
 
 from ._assistant import Assistant, Message, MessageRole
 
-from ._component import RagComponent
 from ._config import Config
-from ._core import RagnaException, RagnaId
+from ._core import RagnaException
 from ._document import Document
 from ._queue import Queue
-from ._source_storage import ReconstructedSource, SourceStorage
-from ._state import State
-
-T = TypeVar("T", bound=RagComponent)
+from ._source_storage import SourceStorage
 
 
 class Rag:
@@ -30,178 +29,43 @@ def __init__(
         self.config = config or Config()
         self._logger = self.config.get_logger()
 
-        self._state = State(self.config.state_database_url)
         self._queue = Queue(self.config, load_components=load_components)
 
-        self._chats: dict[(str, str), Chat] = {}
-
-    async def new_chat(
+    def chat(
         self,
         *,
-        user: str = "Ragna",
-        name: Optional[str] = None,
-        documents: Sequence[Any],
-        source_storage: Union[Type[RagComponent], RagComponent, str],
-        assistant: Union[Type[RagComponent], RagComponent, str],
-        **params,
+        documents: Iterable[Any],
+        source_storage: Union[Type[SourceStorage], SourceStorage, str],
+        assistant: Union[Type[Assistant], Assistant, str],
+        **params: Any,
     ):
-        """Create a new [ragna.core.Chat][] object from the given configuration."""
-        documents = self._parse_documents(documents, user=user)
-        source_storage = self._queue.parse_component(source_storage, load=True)
-        assistant = self._queue.parse_component(assistant, load=True)
-
-        chat = Chat(
-            rag=self,
-            user=user,
-            id=RagnaId.make(),
-            name=name,
+        """Create a new [ragna.core.Chat][]."""
+        return Chat(
+            self,
             documents=documents,
             source_storage=source_storage,
             assistant=assistant,
             **params,
         )
 
-        self._state.add_chat(
-            id=chat.id,
-            user=user,
-            name=chat.name,
-            document_ids=[document.id for document in documents],
-            source_storage=source_storage.display_name(),
-            assistant=assistant.display_name(),
-            params=params,
-        )
-
-        return chat
-
-    def _parse_documents(
-        self, documents: Sequence[Any], *, user: str
-    ) -> list[Document]:
-        documents_ = []
-        for document in documents:
-            if isinstance(document, RagnaId):
-                document = self._get_document(id=document, user=user)
-            else:
-                if not isinstance(document, Document):
-                    document = self.config.document_class(document)
-
-                if document.id is None:
-                    document.id = RagnaId.make()
-                    self._add_document(
-                        user=user,
-                        id=document.id,
-                        name=document.name,
-                        metadata=document.metadata,
-                    )
-
-            if not document.is_available():
-                raise RagnaException(
-                    "Document not available",
-                    document=document,
-                    http_status_code=404,
-                    http_detail=f"Document with ID {document.id} not available.",
-                )
-
-            documents_.append(document)
-        return documents_
-
-    def _add_document(self, *, user: str, id: RagnaId, name: str, metadata):
-        self._state.add_document(user=user, id=id, name=name, metadata=metadata)
-
-    def _get_document(self, user: str, id: RagnaId):
-        state = self._state.get_document(user=user, id=id)
-        if state is None:
-            raise RagnaException(
-                "Document not found",
-                user=user,
-                id=id,
-                http_status_code=404,
-                http_detail=RagnaException.EVENT,
-            )
-        return self.config.document_class(
-            id=id, name=state.name, metadata=state.metadata_
-        )
-
-    def _get_chats(self, *, user: str):
-        chats = [
-            Chat(
-                rag=self,
-                user=user,
-                id=chat_state.id,
-                name=chat_state.name,
-                documents=[
-                    self.config.document_class(
-                        id=document_state.id,
-                        name=document_state.name,
-                        metadata=document_state.metadata_,
-                    )
-                    for document_state in chat_state.document_states
-                ],
-                source_storage=self._queue.parse_component(chat_state.source_storage),
-                assistant=self._queue.parse_component(chat_state.assistant),
-                messages=[
-                    Message(
-                        id=message_state.id,
-                        content=message_state.content,
-                        role=message_state.role,
-                        sources=[
-                            ReconstructedSource(
-                                id=source_state.id,
-                                document_id=source_state.document_id,
-                                document_name=source_state.document_state.name,
-                                location=source_state.location,
-                            )
-                            for source_state in message_state.source_states
-                        ],
-                    )
-                    for message_state in chat_state.message_states
-                ],
-                **chat_state.params,
-            )
-            for chat_state in self._state.get_chats(user=user)
-        ]
-        self._chats.update({(user, chat.id): chat for chat in chats})
-        return chats
-
-    def _get_chat(self, *, user: str, id: RagnaId):
-        key = (user, id)
-
-        chat = self._chats.get(key)
-        if chat is not None:
-            return chat
-
-        self._get_chats(user=user)
-
-        chat = self._chats.get(key)
-        if chat is not None:
-            return chat
-
-        raise RagnaException(
-            "Chat not found",
-            user=user,
-            id=id,
-            http_status_code=404,
-            detail=RagnaException.EVENT,
-        )
-
 
 class Chat:
     """
     !!! note
 
         This object is usually not instantiated manually, but rather through
-        [ragna.core.Rag.new_chat][].
+        [ragna.core.Rag.chat][].
 
-    A chat needs to be [`start`][ragna.core.Chat.start]ed before prompts can be
-    [`answer`][ragna.core.Chat.answer]ed. Optionally, it can be
-    [`close`][ragna.core.Chat.close]d to no longer accept new prompts.
+    A chat needs to be [`prepare`][ragna.core.Chat.prepare]d before prompts can be
+    [`answer`][ragna.core.Chat.answer]ed.
 
     Can be used as context manager to automatically invoke
-    [`start`][ragna.core.Chat.start] and [`close`][ragna.core.Chat.close]:
+    [`prepare`][ragna.core.Chat.prepare]:
 
     ```python
     rag = Rag()
 
-    async with await rag.new_chat(
+    async with rag.chat(
         documents=[path],
         source_storage=ragna.core.RagnaDemoSourceStorage,
         assistant=ragna.core.RagnaDemoAssistant,
@@ -212,136 +76,111 @@ class Chat:
 
     def __init__(
         self,
-        *,
         rag: Rag,
-        user: str,
-        id: RagnaId,
-        name: Optional[str] = None,
-        documents,
-        source_storage: Type[SourceStorage],
-        assistant: Type[Assistant],
-        messages: Optional[list[Message]] = None,
-        **params,
+        *,
+        documents: Iterable[Any],
+        source_storage: Union[Type[SourceStorage], SourceStorage, str],
+        assistant: Union[Type[Assistant], Assistant, str],
+        **params: Any,
     ):
         self._rag = rag
-        self._state = self._rag._state
-        self._user = user
 
-        self.id = id
-        self.name = name or f"{datetime.datetime.now():%c}"
-        self.documents = documents
-        self.source_storage = source_storage
-        self.assistant = assistant
+        self.documents = self._parse_documents(documents)
+        self.source_storage = self._rag._queue.parse_component(source_storage)
+        self.assistant = self._rag._queue.parse_component(assistant)
 
+        special_params = self._SpecialChatParams().dict()
+        special_params.update(params)
+        params = special_params
         self.params = params
         self._unpacked_params = self._unpack_chat_params(params)
 
-        self.messages = messages or []
+        self._prepared = False
+        self._messages = []
 
-        self._started = False
-        self._closed = False
-
-    async def start(self):
-        """Start the chat.
+    async def prepare(self):
+        """Prepare the chat.
 
         This [`store`][ragna.core.SourceStorage.store]s the documents in the selected
-        source storage.
+        source storage. Afterwards prompts can be [`answer`][ragna.core.Chat.answer]ed.
 
         Raises:
             ragna.core.RagnaException: If chat is already
-                [`start`][ragna.core.Chat.start]'ed.
-            ragna.core.RagnaException: If chat is [`close`][ragna.core.Chat.close]'ed.
+                [`prepare`][ragna.core.Chat.prepare]d.
         """
-        if self._started:
+        if self._prepared:
             raise RagnaException(
-                "Chat is already started",
+                "Chat is already prepared",
                 chat=self,
                 http_status_code=400,
                 detail=RagnaException.EVENT,
             )
-        elif self._closed:
-            raise RagnaException(
-                "Chat is closed and cannot be restarted",
-                chat=self,
-                http_status_code=400,
-                http_detail=RagnaException.EVENT,
-            )
 
         await self._enqueue(self.source_storage, "store", self.documents)
-        self._state.start_chat(user=self._user, id=self.id)
-        self._started = True
+        self._prepared = True
 
         welcome = Message(
-            id=RagnaId.make(),
             content="How can I help you with the documents?",
             role=MessageRole.SYSTEM,
         )
-        self._append_message(welcome)
-
-        return self
-
-    async def close(self):
-        """Close that chat.
-
-        After the chat is closed, new prompts will no longer be
-        [`answer`][ragna.core.Chat.answer]'ed.
-        """
-        self._state.close_chat(id=self.id, user=self._user)
-        self._closed = True
-
-        return self
+        self._messages.append(welcome)
+        return welcome
 
     async def answer(self, prompt: str):
         """Answer a prompt
 
         Raises:
             ragna.core.RagnaException: If chat is not
-                [`start`][ragna.core.Chat.start]'ed.
-            ragna.core.RagnaException: If chat is [`close`][ragna.core.Chat.close]'ed.
+                [`prepare`][ragna.core.Chat.prepare]d.
         """
-
-        if not self._started:
+        if not self._prepared:
             raise RagnaException(
-                "Chat is not started",
+                "Chat is not prepared",
                 chat=self,
                 http_status_code=400,
                 detail=RagnaException.EVENT,
             )
-        elif self._closed:
-            raise RagnaException(
-                "Chat is closed",
-                chat=self,
-                http_status_code=400,
-                http_detail=RagnaException.EVENT,
-            )
 
-        prompt = Message(id=RagnaId.make(), content=prompt, role=MessageRole.USER)
-        self._append_message(prompt)
+        prompt = Message(content=prompt, role=MessageRole.USER)
+        self._messages.append(prompt)
 
-        sources = await self._enqueue(self.source_storage, "retrieve", prompt.content)
-        content = await self._enqueue(self.assistant, "answer", prompt.content, sources)
+        sources = await self._enqueue(
+            self.source_storage, "retrieve", self.documents, prompt.content
+        )
 
         answer = Message(
-            id=RagnaId.make(),
-            content=content,
+            content=await self._enqueue(
+                self.assistant, "answer", prompt.content, sources
+            ),
             role=MessageRole.ASSISTANT,
             sources=sources,
         )
-        self._append_message(answer)
+        self._messages.append(answer)
+
         return answer
 
-    def _append_message(self, message: Message):
-        self.messages.append(message)
-        self._state.add_message(
-            message,
-            user=self._user,
-            chat_id=self.id,
-        )
+    def _parse_documents(self, documents: Iterable[Any]) -> list[Document]:
+        documents_ = []
+        for document in documents:
+            if not isinstance(document, Document):
+                document = self._rag.config.document_class(document)
+
+            if not document.is_available():
+                raise RagnaException(
+                    "Document not available",
+                    document=document,
+                    http_status_code=404,
+                )
+
+            documents_.append(document)
+        return documents_
 
     class _SpecialChatParams(BaseModel):
-        user: str
-        chat_id: RagnaId
-        chat_name: str
+        user: str = Field(default_factory=os.getlogin)
+        chat_id: uuid.UUID = Field(default_factory=uuid.uuid4)
+        chat_name: str = Field(
+            default_factory=lambda: f"Chat {datetime.datetime.now():%x %X}"
+        )
 
     def _unpack_chat_params(self, params):
         source_storage_models = self.source_storage._models()
@@ -353,13 +192,7 @@ def _unpack_chat_params(self, params):
             *assistant_models.values(),
         )
 
-        chat_model = ChatModel(
-            user=self._user,
-            chat_id=self.id,
-            chat_name=self.name,
-            **params,
-        )
-        chat_params = chat_model.dict(exclude_none=True)
+        chat_params = ChatModel(**params).dict(exclude_none=True)
         return {
             component_and_action: model(**chat_params).dict()
             for component_and_action, model in itertools.chain(
@@ -371,24 +204,18 @@ def _merge_models(self, *models):
         raw_field_definitions = defaultdict(list)
         for model_cls in models:
             for name, field in model_cls.__fields__.items():
-                raw_field_definitions[name].append(
-                    (field.type_, ... if field.required else field.default)
-                )
+                raw_field_definitions[name].append((field.type_, field.required))
 
         field_definitions = {}
         for name, definitions in raw_field_definitions.items():
-            if len(definitions) == 1:
-                field_definitions[name] = definitions[0]
-                continue
-
-            types, defaults = zip(*definitions)
+            types, requireds = zip(*definitions)
 
             types = set(types)
             if len(types) > 1:
                 raise RagnaException(f"Mismatching types for field {name}: {types}")
             type_ = types.pop()
 
-            default = ... if any(default is ... for default in defaults) else None
+            default = ... if any(requireds) else None
 
             field_definitions[name] = (type_, default)
 
@@ -408,8 +235,8 @@ async def _enqueue(self, component, action, *args):
             raise exc
 
     async def __aenter__(self):
-        await self.start()
+        await self.prepare()
         return self
 
-    async def __aexit__(self, *_):
-        await self.close()
+    async def __aexit__(self, *exc_info):
+        pass
diff --git a/ragna/core/_source_storage.py b/ragna/core/_source_storage.py
index 927fb62e..53b8be0d 100644
--- a/ragna/core/_source_storage.py
+++ b/ragna/core/_source_storage.py
@@ -1,8 +1,9 @@
 import abc
-from typing import Protocol, Sequence
+from typing import Optional, Protocol, Sequence
+
+from pydantic import BaseModel
 
 from ._component import RagComponent
-from ._core import RagnaException, RagnaId
 from ._document import Document
 
 
@@ -15,33 +16,15 @@ def decode(self, tokens: Sequence[int]) -> str:
         ...
 
 
-class Source:
-    def __init__(
-        self,
-        *,
-        id: RagnaId,
-        document_id: RagnaId,
-        document_name: str,
-        location: str,
-        content: str,
-        num_tokens: int,
-    ):
-        self.id = id
-        self.document_id = document_id
-        self.document_name = document_name
-        self.location = location
-        self.content = content
-        self.num_tokens = num_tokens
-
-
-class ReconstructedSource(Source):
-    def __init__(self, **kwargs):
-        if any(
-            kwargs.setdefault(param, None) is not None
-            for param in ["content", "num_tokens"]
-        ):
-            raise RagnaException
-        super().__init__(**kwargs)
+class Source(BaseModel):
+    class Config:
+        arbitrary_types_allowed = True
+
+    id: str
+    document: Document
+    location: str
+    content: Optional[str]
+    num_tokens: Optional[int]
 
 
 class SourceStorage(RagComponent, abc.ABC):
@@ -57,5 +40,5 @@ def store(self, documents: list[Document]) -> None:
         ...
 
     @abc.abstractmethod
-    def retrieve(self, prompt: str) -> list[Source]:
+    def retrieve(self, documents: list[Document], prompt: str) -> list[Source]:
         ...
diff --git a/ragna/core/_state.py b/ragna/core/_state.py
deleted file mode 100644
index 710795c7..00000000
--- a/ragna/core/_state.py
+++ /dev/null
@@ -1,180 +0,0 @@
-from __future__ import annotations
-
-import functools
-from typing import Any
-
-from sqlalchemy import create_engine, select
-from sqlalchemy.orm import Session
-
-from ragna.core import Message, RagnaException, RagnaId
-
-from ._orm import Base, ChatState, DocumentState, MessageState, SourceState, UserState
-
-
-class State:
-    def __init__(self, url: str):
-        self._engine = create_engine(url)
-        Base.metadata.create_all(self._engine)
-        self._session = Session(self._engine)
-
-    def __del__(self):
-        if hasattr(self, "_session"):
-            self._session.close()
-
-    @functools.lru_cache(maxsize=1024)
-    def _get_user_id(self, user: str):
-        user_state = self._session.execute(
-            select(UserState).where(UserState.name == user)
-        ).scalar_one_or_none()
-        if user_state is not None:
-            return user_state.id
-
-        user_state = UserState(id=RagnaId.make(), name=user)
-        self._session.add(user_state)
-        self._session.commit()
-        return user_state.id
-
-    def add_document(
-        self, *, user: str, id: RagnaId, name: str, metadata: dict[str, Any]
-    ):
-        document_state = DocumentState(
-            id=id,
-            user_id=self._get_user_id(user),
-            name=name,
-            metadata_=metadata,
-        )
-        self._session.add(document_state)
-        self._session.commit()
-        return document_state
-
-    @functools.lru_cache(maxsize=1024)
-    def get_document(self, user: str, id: RagnaId) -> DocumentState:
-        return self._session.execute(
-            select(DocumentState).where(
-                (DocumentState.user_id == self._get_user_id(user))
-                & (DocumentState.id == id)
-            )
-        ).scalar_one_or_none()
-
-    def get_chats(self, user: str):
-        # FIXME: Add filters for started and closed here
-        return (
-            self._session.execute(
-                select(ChatState).where(ChatState.user_id == self._get_user_id(user))
-            )
-            .scalars()
-            .all()
-        )
-
-    def add_chat(
-        self,
-        *,
-        user: str,
-        id: RagnaId,
-        name: str,
-        document_ids: list[RagnaId],
-        source_storage: str,
-        assistant: str,
-        params,
-    ) -> ChatState:
-        document_states = (
-            self._session.execute(
-                select(DocumentState).where(DocumentState.id.in_(document_ids))
-            )
-            .scalars()
-            .all()
-        )
-        if len(document_states) != len(document_ids):
-            raise RagnaException(
-                set(document_ids) - {document.id for document in document_states}
-            )
-        chat = ChatState(
-            id=id,
-            user_id=self._get_user_id(user),
-            name=name,
-            document_states=document_states,
-            source_storage=source_storage,
-            assistant=assistant,
-            params=params,
-            started=False,
-            closed=False,
-        )
-        self._session.add(chat)
-        self._session.commit()
-        return chat
-
-    def get_chat(self, *, user: str, id: RagnaId):
-        chat_state = self._session.execute(
-            select(ChatState).where(
-                (ChatState.id == id) & (ChatState.user_id == self._get_user_id(user))
-            )
-        ).scalar_one_or_none()
-        if chat_state is None:
-            raise RagnaException()
-        return chat_state
-
-    def start_chat(self, *, user: str, id: RagnaId):
-        chat_state = self.get_chat(user=user, id=id)
-        chat_state.started = True
-        self._session.commit()
-
-    def close_chat(self, *, user: str, id: RagnaId):
-        chat_state = self.get_chat(user=user, id=id)
-        chat_state.closed = True
-        self._session.commit()
-
-    def add_message(
-        self,
-        message: Message,
-        *,
-        user: str,
-        chat_id: RagnaId,
-    ):
-        chat_state = self._session.execute(
-            select(ChatState).where(
-                (ChatState.user_id == self._get_user_id(user))
-                & (ChatState.id == chat_id)
-            )
-        ).scalar_one_or_none()
-        if chat_state is None:
-            raise RagnaException
-
-        if message.sources is not None:
-            sources = {s.id: s for s in message.sources}
-            source_states = list(
-                self._session.execute(
-                    select(SourceState).where(SourceState.id.in_(sources.keys()))
-                )
-                .scalars()
-                .all()
-            )
-            missing_source_ids = sources.keys() - {state.id for state in source_states}
-            if missing_source_ids:
-                for id in missing_source_ids:
-                    source = sources[id]
-                    source_state = SourceState(
-                        id=RagnaId.make(),
-                        document_id=source.document_id,
-                        document_state=self.get_document(
-                            user=user, id=source.document_id
-                        ),
-                        location=source.location,
-                    )
-                    self._session.add(source_state)
-                    source_states.append(source_state)
-        else:
-            source_states = []
-
-        message_state = MessageState(
-            id=message.id,
-            chat_id=chat_state.id,
-            content=message.content,
-            role=message.role,
-            source_states=source_states,
-            timestamp=message.timestamp,
-        )
-        self._session.add(message_state)
-
-        chat_state.message_states.append(message_state)
-
-        self._session.commit()
diff --git a/ragna/source_storage/_chroma.py b/ragna/source_storage/_chroma.py
index 3f674f97..2ecd5854 100644
--- a/ragna/source_storage/_chroma.py
+++ b/ragna/source_storage/_chroma.py
@@ -1,12 +1,6 @@
-from ragna.core import (
-    Document,
-    PackageRequirement,
-    RagnaId,
-    Requirement,
-    Source,
-    SourceStorage,
-)
+import uuid
 
+from ragna.core import Document, PackageRequirement, Requirement, Source, SourceStorage
 from ragna.utils import chunk_pages, page_numbers_to_str, take_sources_up_to_max_tokens
 
 
@@ -45,7 +39,7 @@ def store(
         self,
         documents: list[Document],
         *,
-        chat_id: RagnaId,
+        chat_id: uuid.UUID,
         chunk_size: int = 500,
         chunk_overlap: int = 250,
     ) -> None:
@@ -63,11 +57,11 @@ def store(
                 chunk_overlap=chunk_overlap,
                 tokenizer=self._tokenizer,
             ):
-                ids.append(str(document.id))
+                ids.append(str(uuid.uuid4()))
                 texts.append(chunk.text)
                 metadatas.append(
                     {
-                        "document_name": document.name,
+                        "document_id": str(document.id),
                         "page_numbers": page_numbers_to_str(chunk.page_numbers),
                         "num_tokens": chunk.num_tokens,
                     }
@@ -77,9 +71,10 @@ def store(
 
     def retrieve(
         self,
+        documents: list[Document],
         prompt: str,
         *,
-        chat_id: RagnaId,
+        chat_id: uuid.UUID,
         chunk_size: int = 500,
         num_tokens: int = 1024,
     ) -> list[Source]:
@@ -120,13 +115,13 @@ def retrieve(
         #  2. Whatever threshold we use is very much dependent on the encoding method
         #  Thus, we likely need to have a callable parameter for this class
 
+        document_map = {str(document.id): document for document in documents}
         return list(
             take_sources_up_to_max_tokens(
                 (
                     Source(
-                        id=RagnaId.make(),
-                        document_id=RagnaId(result["ids"]),
-                        document_name=result["metadatas"]["document_name"],
+                        id=result["ids"],
+                        document=document_map[result["metadatas"]["document_id"]],
                         location=result["metadatas"]["page_numbers"],
                         content=result["documents"],
                         num_tokens=result["metadatas"]["num_tokens"],
diff --git a/ragna/source_storage/_demo.py b/ragna/source_storage/_demo.py
index 7073424e..14841405 100644
--- a/ragna/source_storage/_demo.py
+++ b/ragna/source_storage/_demo.py
@@ -1,7 +1,8 @@
 import textwrap
+import uuid
 from typing import Any
 
-from ragna.core import Document, RagnaId, Source, SourceStorage
+from ragna.core import Document, Source, SourceStorage
 
 
 class RagnaDemoSourceStorage(SourceStorage):
@@ -11,31 +12,23 @@ def display_name(cls):
 
     def __init__(self, config):
         super().__init__(config)
-        self._storage: dict[RagnaId, Any] = {}
+        self._storage: dict[uuid.UUID, Any] = {}
 
-    def store(self, documents: list[Document], *, chat_id: RagnaId) -> None:
+    def store(self, documents: list[Document], *, chat_id: uuid.UUID) -> None:
         self._storage[chat_id] = [
-            {
-                "document_id": str(document.id),
-                "document_name": document.name,
-                "location": f"page {page.number}"
+            Source(
+                id=str(uuid.uuid4()),
+                document=document,
+                location=f"page {page.number}"
                 if (page := next(document.extract_pages())).number
                 else "",
-                "content": (content := textwrap.shorten(page.text, width=100)),
-                "num_tokens": len(content.split()),
-            }
+                content=(content := textwrap.shorten(page.text, width=100)),
+                num_tokens=len(content.split()),
+            )
             for document in documents
         ]
 
-    def retrieve(self, prompt: str, *, chat_id: RagnaId) -> list[Source]:
-        return [
-            Source(
-                id=RagnaId.make(),
-                document_id=RagnaId(source["document_id"]),
-                document_name=source["document_name"],
-                location=source["location"],
-                content=source["content"],
-                num_tokens=source["num_tokens"],
-            )
-            for source in self._storage[chat_id]
-        ]
+    def retrieve(
+        self, documents: list[Document], prompt: str, *, chat_id: uuid.UUID
+    ) -> list[Source]:
+        return self._storage[chat_id]
diff --git a/ragna/source_storage/_lancedb.py b/ragna/source_storage/_lancedb.py
index 7d5789e9..c03b5990 100644
--- a/ragna/source_storage/_lancedb.py
+++ b/ragna/source_storage/_lancedb.py
@@ -1,12 +1,6 @@
-from ragna.core import (
-    Document,
-    PackageRequirement,
-    RagnaId,
-    Requirement,
-    Source,
-    SourceStorage,
-)
+import uuid
 
+from ragna.core import Document, PackageRequirement, Requirement, Source, SourceStorage
 from ragna.utils import chunk_pages, page_numbers_to_str, take_sources_up_to_max_tokens
 
 
@@ -36,8 +30,8 @@ def __init__(self, config):
         self._model = SentenceTransformer("paraphrase-albert-small-v2")
         self._schema = pa.schema(
             [
+                pa.field("id", pa.string()),
                 pa.field("document_id", pa.string()),
-                pa.field("document_name", pa.string()),
                 pa.field("page_numbers", pa.string()),
                 pa.field("text", pa.string()),
                 pa.field(
@@ -57,7 +51,7 @@ def store(
         self,
         documents: list[Document],
         *,
-        chat_id: RagnaId,
+        chat_id: uuid.UUID,
         chunk_size: int = 500,
         chunk_overlap: int = 250,
     ) -> None:
@@ -73,8 +67,8 @@ def store(
                 table.add(
                     [
                         {
+                            "id": str(uuid.uuid4()),
                             "document_id": str(document.id),
-                            "document_name": document.name,
                             "page_numbers": page_numbers_to_str(chunk.page_numbers),
                             "text": chunk.text,
                             self._VECTOR_COLUMN_NAME: self._model.encode(chunk.text),
@@ -85,9 +79,10 @@ def store(
 
     def retrieve(
         self,
+        documents: list[Document],
         prompt: str,
         *,
-        chat_id: RagnaId,
+        chat_id: uuid.UUID,
         chunk_size: int = 500,
         num_tokens: int = 1024,
     ) -> list[Source]:
@@ -97,16 +92,23 @@ def retrieve(
         # many sources we have to query. We overestimate by a factor of two to avoid
         # retrieving to few sources and needed to query again.
         limit = int(num_tokens * 2 / chunk_size)
-        results = table.search().limit(limit).to_arrow()
+        results = (
+            table.search(vector_column_name=self._VECTOR_COLUMN_NAME)
+            .limit(limit)
+            .to_arrow()
+        )
 
+        document_map = {str(document.id): document for document in documents}
         return list(
             take_sources_up_to_max_tokens(
                 (
                     Source(
-                        id=RagnaId.make(),
-                        document_id=RagnaId(result["document_id"]),
-                        document_name=result["document_name"],
-                        location=result["page_numbers"],
+                        id=result["id"],
+                        document=document_map[result["document_id"]],
+                        # For some reason adding an empty string during store() results
+                        # in this field being None. Thus, we need to parse it back here.
+                        # TODO: See if there is a configuration option for this
+                        location=result["page_numbers"] or "",
                         content=result["text"],
                         num_tokens=result["num_tokens"],
                     )
diff --git a/scripts/add_chats.py b/scripts/add_chats.py
index 29e358bd..ac8f6994 100644
--- a/scripts/add_chats.py
+++ b/scripts/add_chats.py
@@ -1,6 +1,5 @@
 import datetime
-
-from pprint import pprint
+import json
 
 import httpx
 
@@ -10,11 +9,11 @@ def main():
     url = "http://127.0.0.1:31476"
     user = "Ragna"
 
-    assert client.get(f"{url}/health").is_success
+    assert client.get(f"{url}").is_success
 
     ## documents
 
-    document_ids = []
+    documents = []
     for i in range(5):
         name = f"document{i}.txt"
         document_info = client.get(
@@ -25,7 +24,7 @@ def main():
             data=document_info["data"],
             files={"file": f"Content of {name}".encode()},
         )
-        document_ids.append(document_info["document"]["id"])
+        documents.append(document_info["document"])
 
     ## chat 1
 
@@ -34,14 +33,14 @@ def main():
         params={"user": user},
         json={
             "name": "Test chat",
-            "document_ids": document_ids[:2],
+            "documents": documents[:2],
             "source_storage": "Ragna/DemoSourceStorage",
             "assistant": "Ragna/DemoAssistant",
             "params": {},
         },
     ).json()
     client.post(
-        f"{url}/chats/{chat['id']}/start",
+        f"{url}/chats/{chat['id']}/prepare",
         params={"user": user},
     )
     client.post(
@@ -56,14 +55,14 @@ def main():
         params={"user": user},
         json={
             "name": f"Chat {datetime.datetime.now():%x %X}",
-            "document_ids": document_ids[2:4],
+            "documents": documents[2:4],
             "source_storage": "Ragna/DemoSourceStorage",
             "assistant": "Ragna/DemoAssistant",
             "params": {},
         },
     ).json()
     client.post(
-        f"{url}/chats/{chat['id']}/start",
+        f"{url}/chats/{chat['id']}/prepare",
         params={"user": user},
     )
     for _ in range(3):
@@ -79,14 +78,14 @@ def main():
         params={"user": user},
         json={
             "name": "Really long chat name that likely needs to be truncated somehow. If you can read this, truncating failed :boom:",
-            "document_ids": [document_ids[i] for i in [0, 2, 4]],
+            "documents": [documents[i] for i in [0, 2, 4]],
             "source_storage": "Ragna/DemoSourceStorage",
             "assistant": "Ragna/DemoAssistant",
             "params": {},
         },
     ).json()
     client.post(
-        f"{url}/chats/{chat['id']}/start",
+        f"{url}/chats/{chat['id']}/prepare",
         params={"user": user},
     )
     client.post(
@@ -101,8 +100,8 @@ def main():
         },
     )
 
-    chats = client.get(f"{url}/chats", params={"user": user}).json()
-    pprint(chats)
+    response = client.get(f"{url}/chats", params={"user": user})
+    print(json.dumps(response.json()))
 
 
 if __name__ == "__main__":