Spaces:
Sleeping
Sleeping
File size: 3,372 Bytes
5623f53 |
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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
from llmops.openai_utils.chatmodel import ChatOpenAI
from llmops.vectordatabase import VectorDatabase
from llmops.openai_utils.prompts import (
UserRolePrompt,
SystemRolePrompt,
AssistantRolePrompt,
)
import datetime
from wandb.sdk.data_types.trace_tree import Trace
class RetrievalAugmentedQAPipeline:
"""
"""
def __init__(self, llm: ChatOpenAI(), vector_db_retriever: VectorDatabase) -> None:
self.llm = llm
self.vector_db_retriever = vector_db_retriever
def run_pipeline(self, user_query: str, raqa_prompt:SystemRolePrompt, user_prompt:UserRolePrompt) -> str:
context_list = self.vector_db_retriever.search_by_text(user_query, k=4)
context_prompt = ""
for context in context_list:
context_prompt += context[0] + "\n"
formatted_system_prompt = raqa_prompt.create_message(context=context_prompt)
formatted_user_prompt = user_prompt.create_message(user_query=user_query)
return self.llm.run([formatted_system_prompt, formatted_user_prompt])
class WandB_RetrievalAugmentedQAPipeline:
def __init__(self, llm: ChatOpenAI(), vector_db_retriever: VectorDatabase, wandb_project = None) -> None:
self.llm = llm
self.vector_db_retriever = vector_db_retriever
self.wandb_project = wandb_project
def run_pipeline(self, user_query: str, raqa_prompt:SystemRolePrompt, user_prompt:UserRolePrompt) -> str:
context_list = self.vector_db_retriever.search_by_text(user_query, k=4)
context_prompt = ""
for context in context_list:
context_prompt += context[0] + "\n"
formatted_system_prompt = raqa_prompt.create_message(context=context_prompt)
formatted_user_prompt = user_prompt.create_message(user_query=user_query)
start_time = datetime.datetime.now().timestamp() * 1000
try:
openai_response = self.llm.run([formatted_system_prompt, formatted_user_prompt], text_only=False)
end_time = datetime.datetime.now().timestamp() * 1000
status = "success"
status_message = (None, )
response_text = openai_response.choices[0].message.content
token_usage = openai_response["usage"].to_dict()
model = openai_response["model"]
except Exception as e:
end_time = datetime.datetime.now().timestamp() * 1000
status = "error"
status_message = str(e)
response_text = ""
token_usage = {}
model = ""
if self.wandb_project:
root_span = Trace(
name="root_span",
kind="llm",
status_code=status,
status_message=status_message,
start_time_ms=start_time,
end_time_ms=end_time,
metadata={
"token_usage" : token_usage,
"model_name" : model
},
inputs= {"system_prompt" : formatted_system_prompt, "user_prompt" : formatted_user_prompt},
outputs= {"response" : response_text}
)
root_span.log(name="openai_trace")
return response_text if response_text else "We ran into an error. Please try again later. Full Error Message: " + status_message
|