Skip to content

Commit 51eff78

Browse files
committed
feat(ollama): use chat api to leverage prompt templates
closes #34
1 parent a04effc commit 51eff78

File tree

4 files changed

+37
-42
lines changed

4 files changed

+37
-42
lines changed

gswikichat/api.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,6 @@ async def api(query, top_k=3, lang='en'):
5454
logger.debug(f'{answer=}')
5555

5656
return {
57-
"answer": answer.data,
57+
"answer": answer.data.content,
5858
"sources": sources
5959
}

gswikichat/llm_config.py

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import os
2-
from haystack_integrations.components.generators.ollama import OllamaGenerator
2+
from haystack_integrations.components.generators.ollama import OllamaChatGenerator
33

44
from .logger import get_logger
55

@@ -8,17 +8,16 @@
88

99
OLLAMA_MODEL_NAME = os.environ.get("OLLAMA_MODEL_NAME")
1010
OLLAMA_URL = os.environ.get("OLLAMA_URL")
11-
OLLAMA_GENERATE_URL = f"{OLLAMA_URL}/api/generate"
11+
OLLAMA_CHAT_URL = f"{OLLAMA_URL}/api/chat"
1212

1313
logger.info(f'Using {OLLAMA_MODEL_NAME=}')
1414
logger.info(f'Endpoint: {OLLAMA_URL=}')
15-
logger.info(f'Generate: {OLLAMA_GENERATE_URL=}')
16-
17-
logger.debug(f'I AM HERE')
15+
logger.info(f'Generate: {OLLAMA_CHAT_URL=}')
1816

1917
logger.info(f"Setting up ollama with {OLLAMA_MODEL_NAME}")
2018

21-
llm = OllamaGenerator(
19+
llm = OllamaChatGenerator(
2220
model=OLLAMA_MODEL_NAME,
23-
url=OLLAMA_GENERATE_URL
21+
url=OLLAMA_CHAT_URL,
22+
timeout=120
2423
)

gswikichat/prompt.py

Lines changed: 7 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,38 +1,28 @@
11
from haystack.components.builders.prompt_builder import PromptBuilder
22

33
prompt_template_en = """
4-
<|system|>
5-
You are a helpful assistant. You answer questions based on the given documents.
6-
Answer based on the documents only. If the information is not in the documents,
7-
say that you cannot find the information.
8-
<|endoftext|>
9-
<|user|>
104
Documents:
115
{% for doc_ in documents %}
126
{{ doc_.content }}
137
{% endfor %}
148
With this documents, answer the following question: {{question}}
15-
<|endoftext|>
16-
<|assistant|>
179
"""
1810

1911
prompt_template_de = """
20-
<|system|>
21-
Du bist ein hilfreicher Assistent. Du beantwortest Fragen basierend auf den vorliegenden Dokumenten.
22-
Beantworte basierend auf den Dokumenten nur. Wenn die Information nicht in den Dokumenten ist,
23-
sage, dass du sie nicht finden kannst.
24-
<|endoftext|>
25-
<|user|>
2612
Dokumente:
2713
{% for doc_ in documents %}
2814
{{ doc_.content }}
2915
{% endfor %}
3016
Mit diesen Dokumenten, beantworte die folgende Frage: {{question}}
31-
<|endoftext|>
32-
<|assistant|>
3317
"""
3418

35-
prompt_builders = {
19+
system_prompts = {
20+
'en': 'You are a helpful assistant. You answer questions based on the given documents. Answer based on the documents only. If the information is not in the documents, say that you cannot find the information.',
21+
'de': 'Du bist ein hilfreicher Assistent. Du beantwortest Fragen basierend auf den vorliegenden Dokumenten. Beantworte basierend auf den Dokumenten nur. Wenn die Information nicht in den Dokumenten ist, sage, dass du sie nicht finden kannst.',
22+
}
23+
24+
user_prompt_builders = {
3625
'en': PromptBuilder(template=prompt_template_en),
3726
'de': PromptBuilder(template=prompt_template_de),
3827
}
28+

gswikichat/rag.py

Lines changed: 23 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,21 @@
11
# from haystack import Pipeline
22
from haystack import Document
33
from haystack.components.builders.answer_builder import AnswerBuilder
4+
from haystack.dataclasses import ChatMessage
45

56
from .llm_config import llm
67
from .logger import get_logger
7-
from .prompt import prompt_builders
8+
from .prompt import user_prompt_builders, system_prompts
89
from .vector_store_interface import embedder, retriever, input_documents
910

1011
# Create logger instance from base logger config in `logger.py`
1112
logger = get_logger(__name__)
1213

1314

14-
def rag_pipeline(query: str = None, top_k: int = 3, lang: str = 'de'):
15+
def rag_pipeline(query: str, top_k: int = 3, lang: str = 'de'):
1516

16-
assert (query is not None)
17-
18-
if isinstance(query, str):
19-
query = Document(content=query)
20-
21-
assert (isinstance(query, Document))
22-
23-
query_embedded = embedder.run([query])
17+
query_document = Document(content=query)
18+
query_embedded = embedder.run([query_document])
2419
query_embedding = query_embedded['documents'][0].embedding
2520

2621
retriever_results = retriever.run(
@@ -35,24 +30,35 @@ def rag_pipeline(query: str = None, top_k: int = 3, lang: str = 'de'):
3530
for retriever_result_ in retriever_results:
3631
logger.debug(retriever_result_)
3732

38-
prompt_builder = prompt_builders[lang]
33+
system_prompt = system_prompts[lang]
34+
user_prompt_builder = user_prompt_builders[lang]
3935

40-
prompt_build = prompt_builder.run(
41-
question=query.content, # As a Document instance, .content returns a string
36+
user_prompt_build = user_prompt_builder.run(
37+
question=query_document.content,
4238
documents=retriever_results['documents']
4339
)
4440

45-
prompt = prompt_build['prompt']
41+
prompt = user_prompt_build['prompt']
4642

4743
logger.debug(f'{prompt=}')
4844

49-
response = llm.run(prompt=prompt, generation_kwargs=None)
45+
messages = [
46+
ChatMessage.from_system(system_prompt),
47+
ChatMessage.from_user(prompt),
48+
]
49+
50+
response = llm.run(
51+
messages,
52+
# generation_kwargs={"temperature": 0.2}
53+
)
54+
55+
logger.debug(response)
5056

5157
answer_builder = AnswerBuilder()
5258
answer_build = answer_builder.run(
53-
query=query.content, # As a Document class, .content returns the string
59+
query=query_document.content,
5460
replies=response['replies'],
55-
meta=response['meta'],
61+
meta=[r.meta for r in response['replies']],
5662
documents=retriever_results['documents'],
5763
pattern=None,
5864
reference_pattern=None

0 commit comments

Comments
 (0)