@@ -83,7 +83,7 @@ async def ollama_stream_generator( # noqa: C901
83
83
class OllamaShim (BaseCompletionHandler ):
84
84
85
85
def __init__ (self , base_url ):
86
- self .client = AsyncClient (host = base_url , timeout = 300 )
86
+ self .client = AsyncClient (host = base_url , timeout = 30 )
87
87
88
88
async def execute_completion (
89
89
self ,
@@ -93,31 +93,35 @@ async def execute_completion(
93
93
is_fim_request : bool = False ,
94
94
) -> Union [ChatResponse , GenerateResponse ]:
95
95
"""Stream response directly from Ollama API."""
96
- if is_fim_request :
97
- prompt = ""
98
- for i in reversed (range (len (request ["messages" ]))):
99
- if request ["messages" ][i ]["role" ] == "user" :
100
- prompt = request ["messages" ][i ]["content" ] # type: ignore
101
- break
102
- if not prompt :
103
- raise ValueError ("No user message found in FIM request" )
96
+ try :
97
+ if is_fim_request :
98
+ prompt = ""
99
+ for i in reversed (range (len (request ["messages" ]))):
100
+ if request ["messages" ][i ]["role" ] == "user" :
101
+ prompt = request ["messages" ][i ]["content" ] # type: ignore
102
+ break
103
+ if not prompt :
104
+ raise ValueError ("No user message found in FIM request" )
104
105
105
- response = await self .client .generate (
106
- model = request ["model" ],
107
- prompt = prompt ,
108
- raw = request .get ("raw" , False ),
109
- suffix = request .get ("suffix" , "" ),
110
- stream = stream ,
111
- options = request ["options" ], # type: ignore
112
- )
113
- else :
114
- response = await self .client .chat (
115
- model = request ["model" ],
116
- messages = request ["messages" ],
117
- stream = stream , # type: ignore
118
- options = request ["options" ], # type: ignore
119
- ) # type: ignore
120
- return response
106
+ response = await self .client .generate (
107
+ model = request ["model" ],
108
+ prompt = prompt ,
109
+ raw = request .get ("raw" , False ),
110
+ suffix = request .get ("suffix" , "" ),
111
+ stream = stream ,
112
+ options = request ["options" ], # type: ignore
113
+ )
114
+ else :
115
+ response = await self .client .chat (
116
+ model = request ["model" ],
117
+ messages = request ["messages" ],
118
+ stream = stream , # type: ignore
119
+ options = request ["options" ], # type: ignore
120
+ ) # type: ignore
121
+ return response
122
+ except Exception as e :
123
+ logger .error (f"Error in Ollama completion: { str (e )} " )
124
+ raise e
121
125
122
126
def _create_streaming_response (
123
127
self ,
0 commit comments