1- import json
21import os
3- from typing import Any , Dict , List
4- import re
2+ from typing import Any
53
64from dotenv import load_dotenv
7- from langchain_core .messages import AIMessage , SystemMessage , HumanMessage
8- from langchain_core .tools import BaseTool
5+ from langchain_core .messages import AIMessage
96from langchain_core .runnables import Runnable
107from langchain_gigachat import GigaChat
118from langchain_ollama import ChatOllama
129from langchain_openai import ChatOpenAI
13- from pydantic import BaseModel , ValidationError
10+ from pydantic import BaseModel
1411
12+ from protollm .connectors .rest_server import ChatRESTServer
1513from protollm .connectors .utils import (get_access_token ,
1614 models_without_function_calling ,
17- models_without_structured_output )
15+ models_without_structured_output ,
16+ generate_system_prompt_with_schema ,
17+ generate_system_prompt_with_tools ,
18+ parse_function_calls ,
19+ parse_custom_structure ,
20+ handle_system_prompt )
1821from protollm .definitions import CONFIG_PATH
1922
2023
@@ -45,23 +48,23 @@ def __init__(self, *args: Any, **kwargs: Any):
4548 def invoke (self , messages : str | list , * args , ** kwargs ) -> AIMessage | dict | BaseModel :
4649
4750 if self ._requires_custom_handling_for_tools () and self ._tools :
48- system_prompt = self ._generate_system_prompt_with_tools ( )
49- messages = self . _handle_system_prompt (messages , system_prompt )
51+ system_prompt = generate_system_prompt_with_tools ( self ._tools , self . _tool_choice_mode )
52+ messages = handle_system_prompt (messages , system_prompt )
5053
5154 if self ._requires_custom_handling_for_structured_output () and self ._response_format :
52- system_prompt = self ._generate_system_prompt_with_schema ( )
53- messages = self . _handle_system_prompt (messages , system_prompt )
55+ system_prompt = generate_system_prompt_with_schema ( self ._response_format )
56+ messages = handle_system_prompt (messages , system_prompt )
5457
5558 response = self ._super_invoke (messages , * args , ** kwargs )
5659
5760 match response :
5861 case AIMessage () if ("<function=" in response .content ):
59- tool_calls = self . _parse_function_calls (response .content )
62+ tool_calls = parse_function_calls (response .content )
6063 if tool_calls :
6164 response .tool_calls = tool_calls
6265 response .content = ""
6366 case AIMessage () if self ._response_format :
64- response = self ._parse_custom_structure ( response )
67+ response = parse_custom_structure ( self ._response_format , response )
6568
6669 return response
6770
@@ -83,87 +86,6 @@ def with_structured_output(self, *args, **kwargs: Any) -> Runnable:
8386 else :
8487 return super ().with_structured_output (* args , ** kwargs )
8588
86- def _generate_system_prompt_with_tools (self ) -> str :
87- """
88- Generates a system prompt with function descriptions and instructions for the model.
89-
90- Returns:
91- System prompt with instructions for calling functions and descriptions of the functions themselves.
92-
93- Raises:
94- ValueError: If tools in an unsupported format have been passed.
95- """
96- tool_descriptions = []
97- match self ._tool_choice_mode :
98- case "auto" | None | "any" | "required" | True :
99- tool_choice_mode = str (self ._tool_choice_mode )
100- case _:
101- tool_choice_mode = f"<<{ self ._tool_choice_mode }
102- for tool in self ._tools :
103- match tool :
104- case dict ():
105- tool_descriptions .append (
106- f"Function name: { tool ['name' ]} \n "
107- f"Description: { tool ['description' ]} \n "
108- f"Parameters: { json .dumps (tool ['parameters' ], ensure_ascii = False )}
109- )
110- case BaseTool ():
111- tool_descriptions .append (
112- f"Function name: { tool .name } \n "
113- f"Description: { tool .description } \n "
114- f"Parameters: { json .dumps (tool .args , ensure_ascii = False )}
115- )
116- case _:
117- raise ValueError (
118- "Unsupported tool type. Try using a dictionary or function with the @tool decorator as tools"
119- )
120- tool_prefix = "You have access to the following functions:\n \n "
121- tool_instructions = (
122- "There are the following 4 function call options:\n "
123- "- str of the form <<tool_name>>: call <<tool_name>> tool.\n "
124- "- 'auto': automatically select a tool (including no tool).\n "
125- "- 'none': don't call a tool.\n "
126- "- 'any' or 'required' or 'True': at least one tool have to be called.\n \n "
127- f"User-selected option - { tool_choice_mode } \n \n "
128- "If you choose to call a function ONLY reply in the following format with no prefix or suffix:\n "
129- '<function=example_function_name>{"example_name": "example_value"}</function>'
130- )
131- return tool_prefix + "\n \n " .join (tool_descriptions ) + "\n \n " + tool_instructions
132-
133- def _generate_system_prompt_with_schema (self ) -> str :
134- """
135- Generates a system prompt with response format descriptions and instructions for the model.
136-
137- Returns:
138- A system prompt with instructions for structured output and descriptions of the response formats themselves.
139-
140- Raises:
141- ValueError: If the structure descriptions for the response were passed in an unsupported format.
142- """
143- schema_descriptions = []
144- match self ._response_format :
145- case list ():
146- schemas = self ._response_format
147- case _:
148- schemas = [self ._response_format ]
149- for schema in schemas :
150- match schema :
151- case dict ():
152- schema_descriptions .append (str (schema ))
153- case _ if issubclass (schema , BaseModel ):
154- schema_descriptions .append (str (schema .model_json_schema ()))
155- case _:
156- raise ValueError (
157- "Unsupported schema type. Try using a description of the answer structure as a dictionary or"
158- " Pydantic model."
159- )
160- schema_prefix = "Generate a JSON object that matches one of the following schemas:\n \n "
161- schema_instructions = (
162- "Your response must contain ONLY valid JSON, parsable by a standard JSON parser. Do not include any"
163- " additional text, explanations, or comments."
164- )
165- return schema_prefix + "\n \n " .join (schema_descriptions ) + "\n \n " + schema_instructions
166-
16789 def _requires_custom_handling_for_tools (self ) -> bool :
16890 """
16991 Determines whether additional processing for tool calling is required for the current model.
@@ -175,84 +97,6 @@ def _requires_custom_handling_for_structured_output(self) -> bool:
17597 Determines whether additional processing for structured output is required for the current model.
17698 """
17799 return any (model_name in self .model_name .lower () for model_name in models_without_structured_output )
178-
179- def _parse_custom_structure (self , response_from_model ) -> dict | BaseModel | None :
180- """
181- Parses the model response into a dictionary or Pydantic class
182-
183- Args:
184- response_from_model: response of a model that does not support structured output by default
185-
186- Raises:
187- ValueError: If a structured response is not obtained
188- """
189- match [self ._response_format ][0 ]:
190- case dict ():
191- try :
192- return json .loads (response_from_model .content )
193- except json .JSONDecodeError as e :
194- raise ValueError (
195- "Failed to return structured output. There may have been a problem with loading JSON from the"
196- f" model.\n { e }
197- )
198- case _ if issubclass ([self ._response_format ][0 ], BaseModel ):
199- for schema in [self ._response_format ]:
200- try :
201- return schema .model_validate_json (response_from_model .content )
202- except ValidationError :
203- continue
204- raise ValueError (
205- "Failed to return structured output. There may have been a problem with validating JSON from the"
206- " model."
207- )
208-
209- @staticmethod
210- def _parse_function_calls (content : str ) -> List [Dict [str , Any ]]:
211- """
212- Parses LLM answer (HTML string) to extract function calls.
213-
214- Args:
215- content: model response as an HTML string
216-
217- Returns:
218- A list of dictionaries in tool_calls format
219-
220- Raises:
221- ValueError: If the arguments for a function call are returned in an incorrect format
222- """
223- tool_calls = []
224- pattern = r"<function=(.*?)>(.*?)</function>"
225- matches = re .findall (pattern , content , re .DOTALL )
226-
227- for match in matches :
228- function_name , function_args = match
229- try :
230- arguments = json .loads (function_args )
231- except json .JSONDecodeError as e :
232- raise ValueError (f"Error when decoding function arguments: { e } )
233-
234- tool_call = {
235- "id" : f"call_{ len (tool_calls ) + 1 } ,
236- "type" : "tool_call" ,
237- "name" : function_name ,
238- "args" : arguments
239- }
240- tool_calls .append (tool_call )
241-
242- return tool_calls
243-
244- @staticmethod
245- def _handle_system_prompt (msgs , sys_prompt ):
246- match msgs :
247- case str ():
248- return [SystemMessage (content = sys_prompt ), HumanMessage (content = msgs )]
249- case list ():
250- if not any (isinstance (msg , SystemMessage ) for msg in msgs ):
251- msgs .insert (0 , SystemMessage (content = sys_prompt ))
252- else :
253- idx = next ((index for index , obj in enumerate (msgs ) if isinstance (obj , SystemMessage )), 0 )
254- msgs [idx ].content += "\n \n " + sys_prompt
255- return msgs
256100
257101
258102def create_llm_connector (model_url : str , * args : Any , ** kwargs : Any ) -> CustomChatOpenAI | GigaChat | ChatOpenAI :
@@ -285,6 +129,9 @@ def create_llm_connector(model_url: str, *args: Any, **kwargs: Any) -> CustomCha
285129 elif "ollama" in model_url :
286130 url_and_name = model_url .split (";" )
287131 return ChatOllama (model = url_and_name [2 ], base_url = url_and_name [1 ], * args , ** kwargs )
132+ elif "self_hosted" in model_url :
133+ url_and_name = model_url .split (";" )
134+ return ChatRESTServer (model = url_and_name [2 ], base_url = url_and_name [1 ], * args , ** kwargs )
288135 elif model_url == "test_model" :
289136 return CustomChatOpenAI (model_name = model_url , api_key = "test" )
290137 else :
0 commit comments