lasse
/
_llm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Enhance ChatResponse handling: add headers support and refactor request method in LLM class

Browse Source
legacy
lasseedfast 5 months ago
parent ba75c7dc97
commit d3f347661b
1 changed files with 88 additions and 49 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Unified View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 137
      _llm/llm.py

137
_llm/llm.py
Unescape View File

@ -4,15 +4,38 @@ import re
import traceback import traceback
from typing import Literal, Optional from typing import Literal, Optional
import tiktoken import tiktoken
import ollama
from ollama import ( from ollama import (
Client, Client,
AsyncClient, AsyncClient,
ResponseError, ResponseError,
ChatResponse,
Options, Options,
) )
from ollama._types import ChatResponse
from pydantic import Field
class ChatResponseWithHeaders(ChatResponse):
headers: dict = Field(default_factory=dict)
def patched_request(self, cls, *args, stream=False, **kwargs):
if stream:
return self._original_request(cls, *args, stream=stream, **kwargs)
raw_response = self._request_raw(*args, **kwargs)
# Use the subclass if cls is ChatResponse
if cls.__name__ == "ChatResponse":
obj = ChatResponseWithHeaders(**raw_response.json())
else:
obj = cls(**raw_response.json())
obj.headers = dict(raw_response.headers)
return obj
ollama.Client._original_request = ollama.Client._request
ollama.Client._request = patched_request
import backoff import backoff
import env_manager import env_manager
try: try:
from colorprinter.colorprinter.print_color import * from colorprinter.colorprinter.print_color import *
except ImportError: except ImportError:
@ -65,6 +88,7 @@ class LLM:
timeout: int = 240, timeout: int = 240,
local_available: bool = False, local_available: bool = False,
on_vpn: bool = False, on_vpn: bool = False,
siltent: bool = False,
) -> None: ) -> None:
""" """
Initialize the assistant with the given parameters. Initialize the assistant with the given parameters.
@ -79,7 +103,7 @@ class LLM:
chosen_backend (str, optional): The backend server to use. If not provided, the least connected server is chosen. chosen_backend (str, optional): The backend server to use. If not provided, the least connected server is chosen.
think (bool): Whether to use thinking mode for reasoning models. Defaults to False. think (bool): Whether to use thinking mode for reasoning models. Defaults to False.
on_vpn (bool): Whether the connection is over VPN and a local path to server can be used. Defaults to False. on_vpn (bool): Whether the connection is over VPN and a local path to server can be used. Defaults to False.
silent (bool): If True, suppresses non-error prints. Defaults to False.
Returns: Returns:
None None
""" """
@ -97,6 +121,7 @@ class LLM:
self.tools = tools or [] self.tools = tools or []
self.local_available = local_available self.local_available = local_available
self.chosen_backend = chosen_backend self.chosen_backend = chosen_backend
self.silent = siltent
headers = { headers = {
"Authorization": f"Basic {self.get_credentials()}", "Authorization": f"Basic {self.get_credentials()}",
@ -111,7 +136,9 @@ class LLM:
self.host_url = f"{os.getenv('LLM_URL')}:{os.getenv('LLM_PORT')}" self.host_url = f"{os.getenv('LLM_URL')}:{os.getenv('LLM_PORT')}"
else: else:
self.host_url = os.getenv("LLM_API_URL").rstrip("/api/chat/") self.host_url = os.getenv("LLM_API_URL").rstrip("/api/chat/")
self.client: Client = Client(host=self.host_url, headers=headers, timeout=timeout) self.client: Client = Client(
host=self.host_url, headers=headers, timeout=timeout
)
self.async_client: AsyncClient = AsyncClient() self.async_client: AsyncClient = AsyncClient()
def get_credentials(self): def get_credentials(self):
@ -147,12 +174,12 @@ class LLM:
self, query, user_input, context, messages, images, model, tools=None self, query, user_input, context, messages, images, model, tools=None
): ):
"""Prepare messages and select the appropriate model, handling images if present.""" """Prepare messages and select the appropriate model, handling images if present."""
if model == "embeddings": if model == "embeddings":
self.messages = [{"role": "user", "content": query}] self.messages = [{"role": "user", "content": query}]
model = self.get_model("embeddings") model = self.get_model("embeddings")
return model return model
if messages: if messages:
messages = [ messages = [
{"role": i["role"], "content": re.sub(r"\s*\n\s*", "\n", i["content"])} {"role": i["role"], "content": re.sub(r"\s*\n\s*", "\n", i["content"])}
@ -180,7 +207,6 @@ class LLM:
]: ]:
model = self.get_model(model) model = self.get_model(model)
self.messages.append(message) self.messages.append(message)
return model return model
@ -227,10 +253,11 @@ class LLM:
"""Call the remote Ollama API synchronously.""" """Call the remote Ollama API synchronously."""
self.call_model = model self.call_model = model
self.client: Client = Client(host=self.host_url, headers=headers, timeout=300) self.client: Client = Client(host=self.host_url, headers=headers, timeout=300)
if self.on_vpn: if not self.silent:
print_yellow(f"🤖 Generating using {model} (remote, on VPN)...") if self.on_vpn:
else: print_yellow(f"🤖 Generating using {model} (remote, on VPN)...")
print_yellow(f"🤖 Generating using {model} (remote)...") else:
print_yellow(f"🤖 Generating using {model} (remote)...")
# If this is an embeddings model, call the embed endpoint instead of chat. # If this is an embeddings model, call the embed endpoint instead of chat.
if model == self.get_model("embeddings"): if model == self.get_model("embeddings"):
@ -244,7 +271,9 @@ class LLM:
input_text = self.messages[-1].get("content", "") input_text = self.messages[-1].get("content", "")
# Use the embed API (synchronous) # Use the embed API (synchronous)
response = self.client.embed(model=model, input=input_text, keep_alive=3600 * 24 * 7) response = self.client.embed(
model=model, input=input_text, keep_alive=3600 * 24 * 7
)
return response return response
response = self.client.chat( response = self.client.chat(
@ -257,6 +286,10 @@ class LLM:
format=format, format=format,
think=think, think=think,
) )
if hasattr(response, "headers") and "x-chosen-backend" in response.headers:
self.chosen_backend = response.headers["x-chosen-backend"]
print_blue(f"Backend used: {self.chosen_backend}")
self.chosen_backend = response.headers.get("x-chosen-backend", None)
return response return response
@backoff.on_exception( @backoff.on_exception(
@ -267,13 +300,14 @@ class LLM:
base=10, base=10,
on_backoff=lambda details: print_yellow( on_backoff=lambda details: print_yellow(
f"Retrying due to error: {details['exception']}" f"Retrying due to error: {details['exception']}"
) ),
) )
async def _call_remote_api_async( async def _call_remote_api_async(
self, model, tools, stream, options, format, headers, think=False self, model, tools, stream, options, format, headers, think=False
): ):
"""Call the remote Ollama API asynchronously.""" """Call the remote Ollama API asynchronously."""
print_yellow(f"🤖 Generating using {model} (remote, async)...") if not self.silent:
print_yellow(f"🤖 Generating using {model} (remote, async)...")
# If embedding model, use async embed endpoint # If embedding model, use async embed endpoint
if model == self.get_model("embeddings"): if model == self.get_model("embeddings"):
@ -303,8 +337,8 @@ class LLM:
def _call_local_ollama(self, model, stream, temperature, think=False): def _call_local_ollama(self, model, stream, temperature, think=False):
"""Call the local Ollama instance synchronously.""" """Call the local Ollama instance synchronously."""
import ollama import ollama
if not self.silent:
print_yellow(f"🤖 Generating using {model} (local)...") print_yellow(f"🤖 Generating using {model} (local)...")
options = {"temperature": temperature} options = {"temperature": temperature}
if stream: if stream:
response_stream = ollama.chat( response_stream = ollama.chat(
@ -321,18 +355,18 @@ class LLM:
chunk_message = chunk["message"] chunk_message = chunk["message"]
content = chunk_message.get("content", "") content = chunk_message.get("content", "")
thinking = chunk_message.get("thinking", None) thinking = chunk_message.get("thinking", None)
yield type( yield type(
"OllamaResponse", "OllamaResponse",
(), (),
{ {
"message": type( "message": type(
"Message", "Message",
(), (),
{ {
"content": content, "content": content,
"thinking": thinking # Include thinking in stream chunks "thinking": thinking, # Include thinking in stream chunks
} },
), ),
"done": chunk.get("done", False), "done": chunk.get("done", False),
}, },
@ -372,7 +406,7 @@ class LLM:
if not self.chat: if not self.chat:
self.messages = [self.messages[0]] self.messages = [self.messages[0]]
return response_obj.message return response_obj.message
# @backoff.on_exception( # @backoff.on_exception(
# backoff.expo, # backoff.expo,
# (ResponseError, TimeoutError), # (ResponseError, TimeoutError),
@ -388,8 +422,8 @@ class LLM:
"""Call the local Ollama instance asynchronously (using a thread pool).""" """Call the local Ollama instance asynchronously (using a thread pool)."""
import ollama import ollama
import asyncio import asyncio
if not self.silent:
print_yellow(f"🤖 Generating using {model} (local, async)...") print_yellow(f"🤖 Generating using {model} (local, async)...")
options = {"temperature": temperature} options = {"temperature": temperature}
loop = asyncio.get_event_loop() loop = asyncio.get_event_loop()
if stream: if stream:
@ -411,18 +445,18 @@ class LLM:
chunk_message = chunk["message"] chunk_message = chunk["message"]
content = chunk_message.get("content", "") content = chunk_message.get("content", "")
thinking = chunk_message.get("thinking", None) thinking = chunk_message.get("thinking", None)
yield type( yield type(
"OllamaResponse", "OllamaResponse",
(), (),
{ {
"message": type( "message": type(
"Message", "Message",
(), (),
{ {
"content": content, "content": content,
"thinking": thinking # Include thinking in async stream chunks "thinking": thinking, # Include thinking in async stream chunks
} },
), ),
"done": chunk.get("done", False), "done": chunk.get("done", False),
}, },
@ -532,12 +566,13 @@ class LLM:
try: try:
headers = self._build_headers(model) headers = self._build_headers(model)
options = self._get_options(temperature) options = self._get_options(temperature)
# Call Ollama server # Call Ollama server)
response: ChatResponse = self._call_remote_api( response: ChatResponse = self._call_remote_api(
model, tools, stream, options, format, headers, think=think model, tools, stream, options, format, headers, think=think
) )
# If using embeddings model, the response is an embed result (not a ChatResponse). # If using embeddings model, the response is an embed result (not a ChatResponse).
if model == self.get_model("embeddings"): if model == self.get_model("embeddings"):
return response return response
@ -548,21 +583,19 @@ class LLM:
if isinstance(response, ChatResponse): if isinstance(response, ChatResponse):
# With native thinking support, content is already clean # With native thinking support, content is already clean
result = response.message.content.strip('"') result = response.message.content.strip('"')
self.messages.append( self.messages.append({"role": "assistant", "content": result})
{"role": "assistant", "content": result}
)
if not self.chat: if not self.chat:
self.messages = [self.messages[0]] self.messages = [self.messages[0]]
Warning = ("Please use reposen.message.content when ising _llm") Warning = "Please use reposen.message.content when ising _llm"
return response return response
else: else:
return "An error occurred." return "An error occurred."
except Exception as e: except Exception as e:
traceback.print_exc() traceback.print_exc()
if self.local_available: if self.local_available:
try: try:
return self._call_local_ollama(model, stream, temperature, think=think) return self._call_local_ollama(model, stream, temperature, think=think)
@ -631,10 +664,8 @@ class LLM:
if isinstance(response, ChatResponse): if isinstance(response, ChatResponse):
# With native thinking support, content is already clean # With native thinking support, content is already clean
result = response.message.content.strip('"') result = response.message.content.strip('"')
self.messages.append( self.messages.append({"role": "assistant", "content": result})
{"role": "assistant", "content": result}
)
if not self.chat: if not self.chat:
self.messages = [self.messages[0]] self.messages = [self.messages[0]]
@ -761,36 +792,44 @@ class LLM:
if __name__ == "__main__": if __name__ == "__main__":
# Example usage of the LLM class with thinking mode # Example usage of the LLM class with thinking mode
llm = LLM() system_message = "You are an extraction assistant. You will get one page at a time from a PDF document, and your task is to extract subcontractors and close clues from the text.\nAssume the state oil company in Uganda is UNOC even if not named. Extract subcontractors (entities or individuals contracted to perform work for UNOC) and close clues.\n\nMANDATES:\n- Do NOT hallucinate. Extract only names or clear clues present in the text.\n- Include \"clues\" such as 'awarded to', 'won the tender', 'appointed', 'supplier', 'contractor', 'consultant', 'tender', 'procured by'.\n- Return EXACT supporting text snippets (<= 300 chars) that justify each extraction.\n- Provide a brief explanation where you explain why the entity is a subcontractor.\n- Returned named entities should be real entities (companies or persons) that could plausibly be subcontractors."
llm = LLM(system_message=system_message)
prompt = "Page 1 of document:\n-----START OF PAGE-----\n**GUIDELINES FOR THE 2019 REGISTRATION ON THE NATIONAL OIL AND GAS**\n**TALENT REGISTER**\n**Welcome to the 2019 National Oil and Gas Talent Register (NOGTR).**\nThe Petroleum Authority of Uganda has developed a National Oil and Gas Talent Register\nto capture all talent that can potentially work in the oil and gas sector as required by law.\nThe NOGTR is a register classified into the demand and supply side users. The demand\nside users consist of companies/government agencies which meet the eligibility criteria\nseeking to recruit human resource across projects in the oil and gas sector and\nredeployment of the same to other sectors.\nOn the other hand, the supply side users include individuals who meet the eligibility criteria\nfor the workforce demands of the oil and gas and are likely to be recruited by the demand\nside users. The NOGTR registration process shall be maintained annually for both the\ndemand and supply side users and the 2019 process will be in line with these guidelines.\nApplicants are encouraged to visit our website and read the submission guidelines carefully before\nregistering.\n**1.** **Annual Registration Calendar**\nThe Authority shall publish the list of entities and persons willing to provide and supply\nlabour force in the oil and gas sector every 31 [st] December of the applicable year. To\nachieve this process the following timelines shall apply.\n**a) Demand side use**\n1. The window for registration shall open on the 1 [st] February 2019 and remain open\nthroughout the year.\n2. The demand side users shall post all available job openings in the Oil and Gas\nSector on the window provided for this purpose by PAU in addition to other\nchannels of advertising that they may opt to use to give wider publicity as required\nby different guidelines.\n3. The demand side users shall have access to the window for purposes of identifying\nand offering employment to any person from the supply side users who meets their\njob description. A considered supply side user shall be contacted and notified of\nthe offer.\n4. Priority shall be given to potential applicants registered on the NOGTR.\n5. The demand side users are encouraged to verify documents uploaded on the\nwindow with the awarding institutions before relying upon them to award jobs.\n6. The Authority shall publish the list of the demand side users that have qualified to\nthe system every 31 [st] December of the applicable year.\n7. The list shall be updated three times; on 31 [st] March 2019, 30 [th] June 2019 and 31 [st]\nAugust 2019 in line with the National Suppliers Database.\n8. The demand side user shall remain on the NOGTR automatically when they\nreapply on the NSD unless if they no longer meet the criteria.\n-----END OF PAGE-----\nPlease extract subcontractors as per the mandates given."
response = llm.generate(query=prompt)
print(response.message.content)
# Basic usage # Basic usage
result = llm.generate( result = llm.generate(
query="I want to add 2 and 2", query="I want to add 2 and 2",
) )
print("Basic result:", result.content) print("Basic result:", result.message.content)
# Example with thinking mode (for reasoning models) # Example with thinking mode (for reasoning models)
print("\n--- Thinking Mode Example ---") print("\n--- Thinking Mode Example ---")
thinking_result = llm.generate( thinking_result = llm.generate(
query="Solve this step by step: If I have 15 apples and give away 7, then buy 3 more, how many do I have?", query="Solve this step by step: If I have 15 apples and give away 7, then buy 3 more, how many do I have?",
model="reasoning", model="reasoning",
think=True think=True,
) )
print("Answer:", thinking_result.content) print("Answer:", thinking_result.message.content)
if hasattr(thinking_result, 'thinking') and thinking_result.thinking: if hasattr(thinking_result, "thinking") and thinking_result.thinking:
print("Model's reasoning:", thinking_result.thinking) print("Model's reasoning:", thinking_result.thinking)
# Example with streaming and thinking # Example with streaming and thinking
print("\n--- Streaming with Thinking Example ---") print("\n--- Streaming with Thinking Example ---")
for chunk_type, chunk_content in llm.generate( for chunk_type, chunk_content in llm.generate(
query="Write a short explanation of photosynthesis", query="Write a short explanation of photosynthesis",
model="reasoning", model="reasoning",
stream=True, stream=True,
think=True think=True,
): ):
if chunk_type == "thinking": if chunk_type == "thinking":
# Use print with blue color escape codes since print_blue doesn't support 'end' parameter # Use print with blue color escape codes since print_blue doesn't support 'end' parameter
print(f"\033[94m {chunk_content}\033[0m", end="") # Show reasoning process in blue print(
f"\033[94m {chunk_content}\033[0m", end=""
) # Show reasoning process in blue
elif chunk_type == "content": elif chunk_type == "content":
print(chunk_content, end="") # Show final answer print(chunk_content, end="") # Show final answer
elif chunk_type == "thinking_complete": elif chunk_type == "thinking_complete":

Write Preview
Loading…
Cancel
Save