Spaces:
Sleeping
Sleeping
File size: 5,788 Bytes
000787f |
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 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 |
import os, asyncio, logging
import configparser
import logging
from dotenv import load_dotenv
# LangChain imports
from langchain_openai import ChatOpenAI
from langchain_anthropic import ChatAnthropic
from langchain_cohere import ChatCohere
from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
from langchain_core.messages import SystemMessage, HumanMessage
# Local .env file
load_dotenv()
def getconfig(configfile_path: str):
"""
Read the config file
Params
----------------
configfile_path: file path of .cfg file
"""
config = configparser.ConfigParser()
try:
config.read_file(open(configfile_path))
return config
except:
logging.warning("config file not found")
# ---------------------------------------------------------------------
# Provider-agnostic authentication and configuration
# ---------------------------------------------------------------------
def get_auth_config(provider: str) -> dict:
"""Get authentication configuration for different providers"""
auth_configs = {
"openai": {"api_key": os.getenv("OPENAI_API_KEY")},
"huggingface": {"api_key": os.getenv("HF_TOKEN")},
"anthropic": {"api_key": os.getenv("ANTHROPIC_API_KEY")},
"cohere": {"api_key": os.getenv("COHERE_API_KEY")},
}
if provider not in auth_configs:
raise ValueError(f"Unsupported provider: {provider}")
auth_config = auth_configs[provider]
api_key = auth_config.get("api_key")
if not api_key:
raise RuntimeError(f"Missing API key for provider '{provider}'. Please set the appropriate environment variable.")
return auth_config
# ---------------------------------------------------------------------
# Model / client initialization
# ---------------------------------------------------------------------
config = getconfig("params.cfg")
PROVIDER = config.get("generator", "PROVIDER")
MODEL = config.get("generator", "MODEL")
MAX_TOKENS = int(config.get("generator", "MAX_TOKENS"))
TEMPERATURE = float(config.get("generator", "TEMPERATURE"))
# Set up authentication for the selected provider
auth_config = get_auth_config(PROVIDER)
def get_chat_model():
"""Initialize the appropriate LangChain chat model based on provider"""
common_params = {
"temperature": TEMPERATURE,
"max_tokens": MAX_TOKENS,
}
if PROVIDER == "openai":
return ChatOpenAI(
model=MODEL,
openai_api_key=auth_config["api_key"],
**common_params
)
elif PROVIDER == "anthropic":
return ChatAnthropic(
model=MODEL,
anthropic_api_key=auth_config["api_key"],
**common_params
)
elif PROVIDER == "cohere":
return ChatCohere(
model=MODEL,
cohere_api_key=auth_config["api_key"],
**common_params
)
elif PROVIDER == "huggingface":
# Initialize HuggingFaceEndpoint with explicit parameters
llm = HuggingFaceEndpoint(
repo_id=MODEL,
huggingfacehub_api_token=auth_config["api_key"],
task="text-generation",
temperature=TEMPERATURE,
max_new_tokens=MAX_TOKENS
)
return ChatHuggingFace(llm=llm)
else:
raise ValueError(f"Unsupported provider: {PROVIDER}")
# Initialize provider-agnostic chat model
chat_model = get_chat_model()
# ---------------------------------------------------------------------
# Core generation function for both Gradio UI and MCP
# ---------------------------------------------------------------------
async def _call_llm(messages: list) -> str:
"""
Provider-agnostic LLM call using LangChain.
Args:
messages: List of LangChain message objects
Returns:
Generated response content as string
"""
try:
# Use async invoke for better performance
response = await chat_model.ainvoke(messages)
return response.content.strip()
except Exception as e:
logging.exception(f"LLM generation failed with provider '{PROVIDER}' and model '{MODEL}': {e}")
raise
def build_messages(question: str, context: str) -> list:
"""
Build messages in LangChain format.
Args:
question: The user's question
context: The relevant context for answering
Returns:
List of LangChain message objects
"""
system_content = (
"You are an expert assistant. Answer the USER question using only the "
"CONTEXT provided. If the context is insufficient say 'I don't know.'"
)
user_content = f"### CONTEXT\n{context}\n\n### USER QUESTION\n{question}"
return [
SystemMessage(content=system_content),
HumanMessage(content=user_content)
]
async def rag_generate(query: str, context: str) -> str:
"""
Generate an answer to a query using provided context through RAG.
This function takes a user query and relevant context, then uses a language model
to generate a comprehensive answer based on the provided information.
Args:
query (str): The user's question or query
context (str): The relevant context/documents to use for answering
Returns:
str: The generated answer based on the query and context
"""
if not query.strip():
return "Error: Query cannot be empty"
if not context.strip():
return "Error: Context cannot be empty"
try:
messages = build_messages(query, context)
answer = await _call_llm(messages)
return answer
except Exception as e:
logging.exception("Generation failed")
return f"Error: {str(e)}" |