File size: 4,891 Bytes
6e1a53e |
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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 |
import asyncio
import json
import logging
import os
from typing import AsyncIterable
from embedchain import App
from embedchain.config import BaseLlmConfig
from langchain.callbacks.streaming_aiter import AsyncIteratorCallbackHandler
from langchain_community.chat_models.huggingface import ChatHuggingFace, HuggingFaceEndpoint
from langchain.schema import HumanMessage, SystemMessage
# App config using OpenAI gpt-3.5-turbo-1106 as LLM
'''
EC_APP_CONFIG = {
"app": {
"config": {
"id": "embedchain-demo-app",
}
},
"llm": {
"provider": "openai",
"config": {
"model": "gpt-3.5-turbo-1106",
}
},
'vectordb': {
'provider': 'chroma',
'config': {
'collection_name': 'rag-full',
'dir': 'db',
'allow_reset': True
}
}
}
'''
# Uncomment this configuration to use Mistral as LLM
EC_APP_CONFIG = {
"app": {
"config": {
"name": "embedchain-opensource-app"
}
},
"llm": {
"provider": "huggingface",
"config": {
"model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
"temperature": 0.1,
"max_tokens": 250,
"top_p": 0.1
}
},
"embedder": {
"provider": "huggingface",
"config": {
"model": "sentence-transformers/all-mpnet-base-v2"
}
},
'vectordb': {
'provider': 'chroma',
'config': {
'collection_name': 'embedchain_store',
'dir': 'db',
'allow_reset': True
}
}
}
async def generate_sources_str(sources_metadata):
"""Generate a string of unique source URLs from the sources metadata."""
seen_urls = set()
unique_sources = [source for source in sources_metadata if source['url'] not in seen_urls and not seen_urls.add(source['url'])]
sources_str = "<sources>\n" + "\n".join(json.dumps(source) for source in unique_sources) + "\n</sources>\n\n"
return sources_str
async def prepare_contexts_for_llm_query(ec_app, query, config, citations):
"""Retrieve contexts from the database and prepare them for the LLM query."""
contexts = ec_app._retrieve_from_database(input_query=query, config=config, where={"app_id": ec_app.config.id}, citations=citations)
if citations and contexts and isinstance(contexts[0], tuple):
return [context[0] for context in contexts]
return contexts
async def generate_messages(ec_app, query, contexts_data_for_llm_query, config):
"""Generate messages to be used in the LLM query."""
messages = []
if config.system_prompt:
messages.append(SystemMessage(content=config.system_prompt))
prompt = ec_app.llm.generate_prompt(query, contexts_data_for_llm_query)
messages.append(HumanMessage(content=prompt))
return messages
async def send_message(query, session_id, number_documents, citations, stream, model) -> AsyncIterable[str]:
ec_app = App.from_config(config=EC_APP_CONFIG)
context = ec_app.search(query, num_documents=number_documents)
sources_str = await generate_sources_str([c['metadata'] for c in context])
ec_app.llm.update_history(app_id=ec_app.config.id, session_id=session_id)
callback = AsyncIteratorCallbackHandler()
#config = BaseLlmConfig(model=model, stream=stream, callbacks=[callback], api_key=os.environ["OPENAI_API_KEY"])
config = BaseLlmConfig(model=model, stream=stream, callbacks=[callback], api_key=os.environ["HUGGINGFACE_ACCESS_TOKEN"])
contexts_data_for_llm_query = await prepare_contexts_for_llm_query(ec_app, query, config, citations)
messages = await generate_messages(ec_app, query, contexts_data_for_llm_query, config)
kwargs = {
"model": model,
"temperature": config.temperature,
"max_tokens": config.max_tokens,
"model_kwargs": {"top_p": config.top_p} if config.top_p else {},
"streaming": stream,
"callbacks": [callback],
"api_key": config.api_key,
"llm" : HuggingFaceEndpoint(
repo_id= model,
temperature= 0.1,
max_new_tokens= 250,
top_p= 0.1,
streaming= stream,
callbacks= [callback],
huggingfacehub_api_token= config.api_key
)
}
llm_task = asyncio.create_task(ChatHuggingFace(**kwargs).agenerate(messages=[messages]))
generated_answer = ""
try:
yield sources_str
async for token in callback.aiter():
yield token
generated_answer += token
except Exception as e:
logging.exception(f"Caught exception: {e}")
finally:
# add conversation in memory
ec_app.llm.add_history(ec_app.config.id, query, generated_answer, session_id=session_id)
callback.done.set()
await llm_task
|