|
1 | 1 | # vim:fenc=utf-8 |
2 | 2 | # |
3 | 3 | # Copyright (C) 2023 dbpunk.com Author imotai <[email protected]> |
4 | | -# |
5 | | -# Licensed under the Apache License, Version 2.0 (the "License"); |
6 | | -# you may not use this file except in compliance with the License. |
7 | | -# You may obtain a copy of the License at |
8 | | -# |
9 | | -# http://www.apache.org/licenses/LICENSE-2.0 |
10 | | -# |
11 | | -# Unless required by applicable law or agreed to in writing, software |
12 | | -# distributed under the License is distributed on an "AS IS" BASIS, |
13 | | -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
14 | | -# See the License for the specific language governing permissions and |
15 | | -# limitations under the License. |
16 | 4 |
|
17 | 5 | import asyncio |
18 | 6 | import logging |
|
32 | 20 | from typing import AsyncIterable, Any, Dict, List, Optional, Sequence, Union, Type |
33 | 21 | from tempfile import gettempdir |
34 | 22 | from grpc.aio import ServicerContext, server |
35 | | -from octopus_kernel.sdk.kernel_sdk import KernelSDK |
| 23 | +from octopus_sdk.kernel_sdk import KernelSDK |
| 24 | +from octopus_sdk.utils import parse_image_filename |
36 | 25 | from .agent_llm import LLMManager |
37 | 26 | from .agent_builder import build_mock_agent, build_openai_agent, build_codellama_agent |
38 | 27 | import databases |
39 | 28 | import orm |
40 | 29 | from datetime import datetime |
41 | | -from .utils import parse_image_filename |
42 | 30 |
|
43 | 31 | config = dotenv_values(".env") |
44 | 32 | LOG_LEVEL = ( |
@@ -72,7 +60,6 @@ class LiteApp(orm.Model): |
72 | 60 | "saved_filenames": orm.String(max_length=512, allow_null=True), |
73 | 61 | } |
74 | 62 |
|
75 | | - |
76 | 63 | class AgentRpcServer(AgentServerServicer): |
77 | 64 |
|
78 | 65 | def __init__(self): |
@@ -164,7 +151,7 @@ async def run( |
164 | 151 | ), |
165 | 152 | ) |
166 | 153 | function_result = None |
167 | | - async for (result, respond) in agent.call_function(lite_app.code): |
| 154 | + async for (result, respond) in agent.call_function(lite_app.code, context): |
168 | 155 | if context.cancelled(): |
169 | 156 | break |
170 | 157 | function_result = result |
@@ -304,6 +291,8 @@ async def worker(task, agent, queue, context): |
304 | 291 | ) |
305 | 292 | yield respond |
306 | 293 | finally: |
| 294 | + is_cancelled = context.cancelled() |
| 295 | + logger.debug(f" the context is cancelled {is_cancelled}") |
307 | 296 | if context.cancelled(): |
308 | 297 | try: |
309 | 298 | logger.warning("cancel the request by stop kernel") |
|
0 commit comments