Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commitd61db7a

Browse files
committed
Merge remote-tracking branch 'origin/development'
2 parentse453add +f7fb56d commitd61db7a

File tree

17 files changed

+806
-58
lines changed

17 files changed

+806
-58
lines changed

‎dive_mcp_host/host/conf/__init__.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,13 @@ class LogConfig(BaseModel):
8282
classOAuthConfig(BaseModel):
8383
"""Config for OAuth."""
8484

85-
redirect_uri:str="dive://mcp.oauth.redirect"
85+
redirect_uri:str|None=None
86+
87+
defget_redirect_uri(self,port:int)->str:
88+
"""Get redirect URI with dynamic port."""
89+
ifself.redirect_uri:
90+
returnself.redirect_uri
91+
returnf"http://localhost:{port}/api/tools/login/oauth/callback"
8692

8793

8894
classEmbedConfig(BaseModel):

‎dive_mcp_host/host/host.py‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -84,12 +84,14 @@ def __init__(
8484
self._config=config
8585
self._model:BaseChatModel|None=None
8686
self._checkpointer:BaseCheckpointSaver[str]|None=None
87+
# Build OAuthManager kwargs, only pass callback_url if redirect_uri is set
88+
oauth_kwargs= {"store":oauth_store}
89+
ifself.config.oauth_config.redirect_uri:
90+
oauth_kwargs["callback_url"]=self.config.oauth_config.redirect_uri
8791
self._tool_manager:ToolManager=ToolManager(
8892
configs=self._config.mcp_servers,
8993
log_config=self.config.log_config,
90-
oauth_manager=OAuthManager(
91-
self.config.oauth_config.redirect_uri,oauth_store
92-
),
94+
oauth_manager=OAuthManager(**oauth_kwargs),
9395
)
9496
self._store=store_manager
9597
self._exit_stack:AsyncExitStack|None=None

‎dive_mcp_host/host/tools/mcp_server.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -728,6 +728,7 @@ def _http_get_client(
728728
timeout=timeout,
729729
httpx_client_factory=self._httpx_client_factory,
730730
sse_read_timeout=sse_read_timeout,
731+
auth=auth,
731732
)
732733
ifself.config.transportin ("streamable"):
733734
returnstreamablehttp_client(
@@ -1013,8 +1014,10 @@ def auth_manager(self) -> OAuthManager:
10131014

10141015
asyncdefcreate_oauth_authorization(self)->AuthorizationProgress:
10151016
"""Authorize the OAuth client."""
1016-
ifself.config.transport!="streamable":
1017-
raiseRuntimeError("Only streamable transport is supported for oauth")
1017+
ifself.config.transport!="streamable"andself.config.transport!="sse":
1018+
raiseRuntimeError(
1019+
"Only streamable and sse transport is supported for oauth"
1020+
)
10181021

10191022
ifself.auth_managerisNone:
10201023
raiseRuntimeError("OAuth manager is not initialized")

‎dive_mcp_host/httpd/_main.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,10 @@ def main() -> None:
7373
report_status_fd=args.report_status_fd,
7474
)
7575

76+
# Set OAuth resource file path if provided
77+
ifargs.oauth_resource_file:
78+
app.set_oauth_resource_file(str(args.oauth_resource_file))
79+
7680
serversocket=socket.socket(
7781
socket.AF_INET6if":"inargs.listenelsesocket.AF_INET,
7882
socket.SOCK_STREAM,

‎dive_mcp_host/httpd/conf/arguments.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,13 @@ class Arguments(BaseModel):
107107
description="Log level to use.",
108108
)
109109

110+
oauth_resource_file:Annotated[StrPath|None,AfterValidator(_convert_path)]= (
111+
Field(
112+
default=None,
113+
description="Custom HTML template for OAuth callback page.",
114+
)
115+
)
116+
110117
@model_validator(mode="after")
111118
defrewrite_default_path(self)->Self:
112119
"""Rewrite default config file path according to working directory."""
Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
"""add_user_token_and_performance_metrics.
2+
3+
Revision ID: 9f8e7d6c5b4a
4+
Revises: 800075668798
5+
Create Date: 2025-11-18 15:52:13.000000
6+
7+
"""
8+
9+
fromcollections.abcimportSequence
10+
11+
importsqlalchemyassa
12+
fromalembicimportop
13+
14+
# revision identifiers, used by Alembic.
15+
revision:str="9f8e7d6c5b4a"
16+
down_revision:str|None="800075668798"
17+
branch_labels:str|Sequence[str]|None=None
18+
depends_on:str|Sequence[str]|None=None
19+
20+
21+
defupgrade()->None:
22+
"""Upgrade schema."""
23+
# ### commands auto generated by Alembic - please adjust! ###
24+
withop.batch_alter_table("resource_usage")asbatch_op:
25+
batch_op.add_column(
26+
sa.Column("user_token",sa.BigInteger(),nullable=False,server_default="0")
27+
)
28+
batch_op.add_column(
29+
sa.Column(
30+
"time_to_first_token",
31+
sa.Float(),
32+
nullable=False,
33+
server_default="0.0",
34+
)
35+
)
36+
batch_op.add_column(
37+
sa.Column(
38+
"tokens_per_second",sa.Float(),nullable=False,server_default="0.0"
39+
)
40+
)
41+
# ### end Alembic commands ###
42+
43+
44+
defdowngrade()->None:
45+
"""Downgrade schema."""
46+
# ### commands auto generated by Alembic - please adjust! ###
47+
withop.batch_alter_table("resource_usage")asbatch_op:
48+
batch_op.drop_column("tokens_per_second")
49+
batch_op.drop_column("time_to_first_token")
50+
batch_op.drop_column("user_token")
51+
# ### end Alembic commands ###

‎dive_mcp_host/httpd/database/models.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,23 @@
1-
fromdatetimeimportdatetime
1+
from __future__importannotations
2+
3+
fromdatetimeimportdatetime# noqa: TC003
24
fromenumimportStrEnum
35

4-
fromlangchain_core.messagesimportToolCall
6+
fromlangchain_core.messagesimportToolCall# noqa: TC002
57
frompydanticimportBaseModel,ConfigDict,Field
68

9+
fromdive_mcp_host.httpd.routers.modelsimportTokenUsage# noqa: TC001
10+
711

812
classResourceUsage(BaseModel):
913
"""Represents information about a language model's usage statistics."""
1014

1115
model:str
1216
total_input_tokens:int
1317
total_output_tokens:int
18+
user_token:int=0
19+
time_to_first_token:float=0.0
20+
tokens_per_second:float=0.0
1421
total_run_time:float
1522

1623

@@ -78,3 +85,4 @@ class ChatMessage(BaseModel):
7885

7986
chat:Chat
8087
messages:list[Message]
88+
token_usage:TokenUsage|None=None

‎dive_mcp_host/httpd/database/msg_store/abstract.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
Message,
99
NewMessage,
1010
QueryInput,
11+
ResourceUsage,
1112
)
1213

1314

@@ -156,3 +157,16 @@ async def get_next_ai_message(
156157
Returns:
157158
Next AI Message object.
158159
"""
160+
161+
@abstractmethod
162+
asyncdefupdate_message_resource_usage(
163+
self,
164+
message_id:str,
165+
resource_usage:"ResourceUsage",
166+
)->None:
167+
"""Update or create resource usage for a message.
168+
169+
Args:
170+
message_id: Unique identifier for the message.
171+
resource_usage: ResourceUsage data to update or create.
172+
"""

‎dive_mcp_host/httpd/database/msg_store/base.py‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -273,6 +273,9 @@ async def create_message(self, message: NewMessage) -> Message:
273273
"model":message.resource_usage.model,
274274
"total_input_tokens":message.resource_usage.total_input_tokens,
275275
"total_output_tokens":message.resource_usage.total_output_tokens,# noqa: E501
276+
"user_token":message.resource_usage.user_token,
277+
"time_to_first_token":message.resource_usage.time_to_first_token,# noqa: E501
278+
"tokens_per_second":message.resource_usage.tokens_per_second,
276279
"total_run_time":message.resource_usage.total_run_time,
277280
},
278281
)
@@ -515,3 +518,50 @@ async def get_next_ai_message(
515518
toolCalls=message.tool_callsor [],
516519
resource_usage=resource_usage,
517520
)
521+
522+
asyncdefupdate_message_resource_usage(
523+
self,
524+
message_id:str,
525+
resource_usage:ResourceUsage,
526+
)->None:
527+
"""Update or create resource usage for a message.
528+
529+
Args:
530+
message_id: Unique identifier for the message.
531+
resource_usage: ResourceUsage data to update or create.
532+
"""
533+
# Check if resource usage already exists for this message
534+
query=select(ORMResourceUsage).where(
535+
ORMResourceUsage.message_id==message_id
536+
)
537+
existing=awaitself._session.scalar(query)
538+
539+
ifexisting:
540+
# Update existing resource usage
541+
update_query= (
542+
update(ORMResourceUsage)
543+
.where(ORMResourceUsage.message_id==message_id)
544+
.values(
545+
model=resource_usage.model,
546+
total_input_tokens=resource_usage.total_input_tokens,
547+
total_output_tokens=resource_usage.total_output_tokens,
548+
user_token=resource_usage.user_token,
549+
time_to_first_token=resource_usage.time_to_first_token,
550+
tokens_per_second=resource_usage.tokens_per_second,
551+
total_run_time=resource_usage.total_run_time,
552+
)
553+
)
554+
awaitself._session.execute(update_query)
555+
else:
556+
# Create new resource usage
557+
insert_query=insert(ORMResourceUsage).values(
558+
message_id=message_id,
559+
model=resource_usage.model,
560+
total_input_tokens=resource_usage.total_input_tokens,
561+
total_output_tokens=resource_usage.total_output_tokens,
562+
user_token=resource_usage.user_token,
563+
time_to_first_token=resource_usage.time_to_first_token,
564+
tokens_per_second=resource_usage.tokens_per_second,
565+
total_run_time=resource_usage.total_run_time,
566+
)
567+
awaitself._session.execute(insert_query)

‎dive_mcp_host/httpd/database/orm_models.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -161,6 +161,9 @@ class ResourceUsage(Base):
161161
model: Model name.
162162
total_input_tokens: Total input tokens.
163163
total_output_tokens: Total output tokens.
164+
user_token: User input tokens.
165+
time_to_first_token: Time to first token in seconds.
166+
tokens_per_second: Tokens per second.
164167
total_run_time: Total run time.
165168
"""
166169

@@ -177,6 +180,9 @@ class ResourceUsage(Base):
177180
model:Mapped[str]=mapped_column(Text())
178181
total_input_tokens:Mapped[int]=mapped_column(BigInteger())
179182
total_output_tokens:Mapped[int]=mapped_column(BigInteger())
183+
user_token:Mapped[int]=mapped_column(BigInteger(),default=0)
184+
time_to_first_token:Mapped[float]=mapped_column(Float(),default=0.0)
185+
tokens_per_second:Mapped[float]=mapped_column(Float(),default=0.0)
180186
total_run_time:Mapped[float]=mapped_column(Float())
181187

182188
message:Mapped["Message"]=relationship(

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp