From bc04e865b03dff2da3b334998569a78a807d49fe Mon Sep 17 00:00:00 2001 From: fatwang2 Date: Fri, 5 Apr 2024 23:31:20 +0800 Subject: [PATCH] Support Claude for beta, without stream mode and history of chat --- requirements.txt | 1 + search4all.py | 268 +++++++++++++++++++++++++++++------------------ 2 files changed, 168 insertions(+), 101 deletions(-) diff --git a/requirements.txt b/requirements.txt index f64975d..5c684e0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,7 @@ requests httpx openai==1.14.1 +anthropic loguru sanic sqlitedict diff --git a/search4all.py b/search4all.py index 246c996..3a83d75 100644 --- a/search4all.py +++ b/search4all.py @@ -5,8 +5,9 @@ import re import requests import traceback import httpx -from typing import Annotated, List, Generator +from typing import AsyncGenerator from openai import AsyncOpenAI +import anthropic from loguru import logger from dotenv import load_dotenv load_dotenv() @@ -379,15 +380,19 @@ def search_with_searchapi(query: str, subscription_key: str): def new_async_client(_app): - return AsyncOpenAI( - api_key=os.getenv("OPENAI_API_KEY"), - base_url=os.getenv("OPENAI_BASE_URL"), - http_client=_app.ctx.http_session, - ) - + if "claude-3" in _app.ctx.model.lower(): + return anthropic.Anthropic( + api_key=os.getenv("ANTHROPIC_API_KEY") + ) + else: + return AsyncOpenAI( + api_key=os.getenv("OPENAI_API_KEY"), + base_url=os.getenv("OPENAI_BASE_URL"), + http_client=_app.ctx.http_session, + ) @app.before_server_start -async def server_init(_app, loop): +async def server_init(_app): """ Initializes global configs. """ @@ -456,40 +461,6 @@ async def server_init(_app, loop): timeout=httpx.Timeout(connect=10, read=120, write=120, pool=10), ) - -async def get_related_questions_bak(_app, query, contexts): - """ - Gets related questions based on the query using the gpt-3.5-turbo model. - """ - - try: - openai_client = new_async_client(_app) - llm_response = await openai_client.chat.completions.create( - model=_app.ctx.model, - messages=[ - {"role": "system", "content": _more_questions_prompt}, - {"role": "user", "content": query}, - ], - max_tokens=1024, - temperature=0.9, - ) - content = llm_response.choices[0].message.content - questions = [] - for line in content.split('\n'): - if line.strip().startswith('1:') or line.strip().startswith('2:') or line.strip().startswith('3:'): - question_text = line.split(':', 1)[1].strip() - questions.append({"question": question_text}) - if questions: - return questions[:5] - else: - return [] - except Exception as e: - logger.error( - "Encountered error while generating related questions:" - f" {e}\n{traceback.format_exc()}" - ) - return [] - async def get_related_questions(_app, query, contexts): """ Gets related questions based on the query and context. @@ -507,14 +478,15 @@ async def get_related_questions(_app, query, contexts): ) try: - openai_client = new_async_client(_app) - tools = [ - { - "type": "function", - "function": { + logger.info('Start getting related questions') + if "claude-3" in _app.ctx.model.lower(): + logger.info('Using Claude-3 model') + client = new_async_client(_app) + tools = [ + { "name": "ask_related_questions", "description": "Get a list of questions related to the original question and context.", - "parameters": { + "input_schema": { "type": "object", "properties": { "questions": { @@ -527,40 +499,93 @@ async def get_related_questions(_app, query, contexts): }, "required": ["questions"] } + } - } - ] - messages=[ - {"role": "system", "content": _lepton_more_questions_prompt}, + ] + response = client.beta.tools.messages.create( + model=_app.ctx.model, + system=_lepton_more_questions_prompt, + max_tokens=4096, + tools=tools, + messages=[ {"role": "user", "content": query}, ] - request_body = { - "model": _app.ctx.model, - "messages": messages, - "max_tokens": 4096, - "tools": tools, - "tool_choice": { - "type": "function", - "function": { - "name": "ask_related_questions" + ) + logger.info('Response received from Claude-3 model') + + if response.content and len(response.content) > 0: + tool_use = response.content[0] + if tool_use.name == "ask_related_questions": # 使用点操作符访问属性 + related = tool_use.input["questions"] # 使用点操作符访问属性 + else: + related = [] + else: + related = [] + if related and isinstance(related, str): + try: + related = json.loads(related) + except json.JSONDecodeError: + logger.error("Failed to parse related questions as JSON") + return [] + logger.info('Successfully got related questions') + return [{"question": question} for question in related[:5]] + else: + logger.info('Using OpenAI model') + openai_client = new_async_client(_app) + tools = [ + { + "type": "function", + "function": { + "name": "ask_related_questions", + "description": "Get a list of questions related to the original question and context.", + "parameters": { + "type": "object", + "properties": { + "questions": { + "type": "array", + "items": { + "type": "string", + "description": "A related question to the original question and context.", + } + } + }, + "required": ["questions"] + } + } + } + ] + messages=[ + {"role": "system", "content": _lepton_more_questions_prompt}, + {"role": "user", "content": query}, + ] + request_body = { + "model": _app.ctx.model, + "messages": messages, + "max_tokens": 4096, + "tools": tools, + "tool_choice": { + "type": "function", + "function": { + "name": "ask_related_questions" + } + }, } - }, - } - llm_response = await openai_client.chat.completions.create(**request_body) - related = llm_response.choices[0].message.tool_calls[0].function.arguments - if isinstance(related, str): - related = json.loads(related) - logger.trace(f"Related questions: {related}") - return [{"question": _} for _ in related["questions"][:5]] + llm_response = await openai_client.chat.completions.create(**request_body) + logger.info('Response received from OpenAI model') + related = llm_response.choices[0].message.tool_calls[0].function.arguments + if isinstance(related, str): + related = json.loads(related) + logger.trace(f"Related questions: {related}") + return [{"question": _} for _ in related["questions"][:5]] + except Exception as e: logger.error( f"Encountered error while generating related questions: {str(e)}" ) return [] - async def _raw_stream_response( _app, contexts, llm_response, related_questions_future -) -> Generator[str, None, None]: +) -> AsyncGenerator[str, None]: """ A generator that yields the raw stream response. You do not need to call this directly. Instead, use the stream_and_upload_to_kv which will also @@ -576,9 +601,16 @@ async def _raw_stream_response( "(The search engine returned nothing for this query. Please take the" " answer with a grain of salt.)\n\n" ) - async for chunk in llm_response: - if chunk.choices: - yield chunk.choices[0].delta.content or "" + + if "claude-3" in _app.ctx.model.lower(): + # Process Claude's stream response + async for text in llm_response: + yield text + else: + # Process OpenAI's stream response + async for chunk in llm_response: + if chunk.choices: + yield chunk.choices[0].delta.content or "" # Third, yield the related questions. If any error happens, we will just # return an empty list. if related_questions_future is not None: @@ -722,42 +754,76 @@ async def query_function(request: sanic.Request): ) ) try: - openai_client = new_async_client(_app) - messages=[ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": query}, - ] - if chat_history and len(chat_history) % 2 == 0: - # 将历史插入到消息中 index = 1 的位置 - messages[1:1] = chat_history - llm_response = await openai_client.chat.completions.create( - model=_app.ctx.model, - messages=messages, - max_tokens=1024, - stream=True, - temperature=0.9, - ) if _app.ctx.should_do_related_questions and generate_related_questions: # While the answer is being generated, we can start generating # related questions as a future. related_questions_future = get_related_questions(_app, query, contexts) - else: related_questions_future = None + if "claude-3" in _app.ctx.model.lower(): + logger.info("Using Claude for generating LLM response") + client = new_async_client(_app) + messages = [ + {"role": "user", "content": query} + ] + if chat_history and len(chat_history) % 2 == 0: + # 确保 chat_history 的第一个消息是一个助手消息 + if chat_history[0]["role"] == "user": + chat_history.insert(0, {"role": "assistant", "content": "okay"}) + messages[1:1] = chat_history + async def _stream_llm_response(stream): + for text in stream.text_stream: + yield text + + # 使用Claude生成LLM响应 + with client.messages.stream( + model=_app.ctx.model, + max_tokens=1024, + system=system_prompt, + messages=messages, + ) as stream: + llm_response = _stream_llm_response(stream) + # 处理流式输出 + + response = await request.respond(content_type="text/html") + # First, stream and yield the results. + all_yielded_results = [] + async for result in _raw_stream_response( + _app, contexts, llm_response, related_questions_future + ): + all_yielded_results.append(result) + await response.send(result) + logger.info("Finished streaming LLM response") + else: + logger.info("Using OpenAI for generating LLM response") + openai_client = new_async_client(_app) + messages=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": query}, + ] + if chat_history and len(chat_history) % 2 == 0: + # 将历史插入到消息中 index = 1 的位置 + messages[1:1] = chat_history + llm_response = await openai_client.chat.completions.create( + model=_app.ctx.model, + messages=messages, + max_tokens=1024, + stream=True, + temperature=0.9, + ) + response = await request.respond(content_type="text/html") + # First, stream and yield the results. + all_yielded_results = [] + async for result in _raw_stream_response( + _app, contexts, llm_response, related_questions_future + ): + all_yielded_results.append(result) + await response.send(result) + logger.info("Finished streaming LLM response") + except Exception as e: logger.error(f"encountered error: {e}\n{traceback.format_exc()}") return sanic.json({"message": "Internal server error."}, 503) - - response = await request.respond(content_type="text/html") - # First, stream and yield the results. - all_yielded_results = [] - - async for result in _raw_stream_response( - _app, contexts, llm_response, related_questions_future - ): - all_yielded_results.append(result) - await response.send(result) - # Second, upload to KV. Note that if uploading to KV fails, we will silently # ignore it, because we don't want to affect the user experience. await response.eof()