-
Notifications
You must be signed in to change notification settings - Fork 0
/
mistralAI.py
33 lines (26 loc) · 1.05 KB
/
mistralAI.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
from mistralai.client import MistralClient
from mistralai.models.chat_completion import ChatMessage
def build_prompt(query, search_results):
prompt_template = """
You're a course teaching assistant. Answer the QUESTION based on the CONTEXT from the FAQ database.
Use only the facts from the CONTEXT when answering the QUESTION.
QUESTION: {question}
CONTEXT:
{context}
""".strip()
context = ""
for doc in search_results:
context = context + f"section: {doc['_source']['section']}\nquestion: {doc['_source']['question']}\nanswer: {doc['_source']['text']}\n\n"
prompt = prompt_template.format(question=query, context=context).strip()
return prompt
def llm(prompt, model="mistral-medium-latest"):
mistral_key = "" # replace with your Mistral.AI KEY
messages = [
ChatMessage(role="user", content=prompt)
]
client = MistralClient(api_key=mistral_key)
chat_response = client.chat(
model=model,
messages=messages
)
return (chat_response.choices[0].message.content)