From 636552ba6efc845b45228029b4c7704b587ff4b0 Mon Sep 17 00:00:00 2001 From: Oyster Oil Date: Tue, 30 Apr 2024 21:14:51 +0800 Subject: [PATCH] Rename files --- .../database/__init__.py => docs/models.rst | 0 ragbear/_entry.py | 36 +++++++++++++++---- ragbear/{database => }/index/__init__.py | 0 .../remote => index/engines}/__init__.py | 0 .../{database/index => index/engines}/hnsw.py | 0 ragbear/{database => index}/find.py | 0 ragbear/{database => index}/ingest.py | 0 ragbear/{database => index}/ingest/extract.py | 0 .../vendor => index/remote}/__init__.py | 0 .../{database => index}/remote/duckduckgo.py | 0 ragbear/{database => index}/schema.py | 0 ragbear/index/vendor/__init__.py | 0 .../{database => index}/vendor/wiki_html.py | 0 ragbear/postprocess/__init__.py | 0 ragbear/{query => synthesis}/prompts.py | 0 15 files changed, 29 insertions(+), 7 deletions(-) rename ragbear/database/__init__.py => docs/models.rst (100%) rename ragbear/{database => }/index/__init__.py (100%) rename ragbear/{database/remote => index/engines}/__init__.py (100%) rename ragbear/{database/index => index/engines}/hnsw.py (100%) rename ragbear/{database => index}/find.py (100%) rename ragbear/{database => index}/ingest.py (100%) rename ragbear/{database => index}/ingest/extract.py (100%) rename ragbear/{database/vendor => index/remote}/__init__.py (100%) rename ragbear/{database => index}/remote/duckduckgo.py (100%) rename ragbear/{database => index}/schema.py (100%) create mode 100644 ragbear/index/vendor/__init__.py rename ragbear/{database => index}/vendor/wiki_html.py (100%) create mode 100644 ragbear/postprocess/__init__.py rename ragbear/{query => synthesis}/prompts.py (100%) diff --git a/ragbear/database/__init__.py b/docs/models.rst similarity index 100% rename from ragbear/database/__init__.py rename to docs/models.rst diff --git a/ragbear/_entry.py b/ragbear/_entry.py index e447fd2..7d30ae6 100644 --- a/ragbear/_entry.py +++ b/ragbear/_entry.py @@ -1,9 +1,9 @@ from contextlib import suppress from typing import List, Literal, Optional, TypedDict, Union -from .query.prompts import format_query -from .model import get_openai_client -from .database.remote.duckduckgo import get_duckduckgo_results +from .index.remote.duckduckgo import get_duckduckgo_results +from .model import get_openai_client +from .synthesis.prompts import format_query with suppress(ImportError): pass @@ -20,6 +20,29 @@ class SourceVendorType(TypedDict): RewriteType = Literal["query2doc"] +class Bear: + + def __init__( + self, + model: str = "llama2", + base_url: str = "http://localhost:11434/v1/", + api_key: Optional[str] = "ollama", + source: Optional[SourceType] = None, + rerank: Optional[RerankType] = None, + rewrite: Optional[RewriteType] = None, + template: Optional[str] = None, + embedding_model: Optional[str] = None, + ): ... + + def answer(self): ... + + def __enter__(self): ... + + def __exit__(self): ... + + def __del__(self): ... + + def query( prompt: str, *, @@ -30,7 +53,7 @@ def query( rerank: Optional[RerankType] = None, rewrite: Optional[RewriteType] = None, template: Optional[str] = None, - embedding_model: Optional[str] = None + embedding_model: Optional[str] = None, ): """Query the LLM, and return the summarized result. @@ -51,12 +74,11 @@ def query( chat_completion = client.chat.completions.create( messages=[ { - 'role': 'user', - 'content': query, + "role": "user", + "content": query, } ], model=model, ) - # print(chat_completion) return chat_completion.choices[0].message.content diff --git a/ragbear/database/index/__init__.py b/ragbear/index/__init__.py similarity index 100% rename from ragbear/database/index/__init__.py rename to ragbear/index/__init__.py diff --git a/ragbear/database/remote/__init__.py b/ragbear/index/engines/__init__.py similarity index 100% rename from ragbear/database/remote/__init__.py rename to ragbear/index/engines/__init__.py diff --git a/ragbear/database/index/hnsw.py b/ragbear/index/engines/hnsw.py similarity index 100% rename from ragbear/database/index/hnsw.py rename to ragbear/index/engines/hnsw.py diff --git a/ragbear/database/find.py b/ragbear/index/find.py similarity index 100% rename from ragbear/database/find.py rename to ragbear/index/find.py diff --git a/ragbear/database/ingest.py b/ragbear/index/ingest.py similarity index 100% rename from ragbear/database/ingest.py rename to ragbear/index/ingest.py diff --git a/ragbear/database/ingest/extract.py b/ragbear/index/ingest/extract.py similarity index 100% rename from ragbear/database/ingest/extract.py rename to ragbear/index/ingest/extract.py diff --git a/ragbear/database/vendor/__init__.py b/ragbear/index/remote/__init__.py similarity index 100% rename from ragbear/database/vendor/__init__.py rename to ragbear/index/remote/__init__.py diff --git a/ragbear/database/remote/duckduckgo.py b/ragbear/index/remote/duckduckgo.py similarity index 100% rename from ragbear/database/remote/duckduckgo.py rename to ragbear/index/remote/duckduckgo.py diff --git a/ragbear/database/schema.py b/ragbear/index/schema.py similarity index 100% rename from ragbear/database/schema.py rename to ragbear/index/schema.py diff --git a/ragbear/index/vendor/__init__.py b/ragbear/index/vendor/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ragbear/database/vendor/wiki_html.py b/ragbear/index/vendor/wiki_html.py similarity index 100% rename from ragbear/database/vendor/wiki_html.py rename to ragbear/index/vendor/wiki_html.py diff --git a/ragbear/postprocess/__init__.py b/ragbear/postprocess/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ragbear/query/prompts.py b/ragbear/synthesis/prompts.py similarity index 100% rename from ragbear/query/prompts.py rename to ragbear/synthesis/prompts.py