From 77eef27cad9b0e5b0551d5a8ea9485779ee7dc18 Mon Sep 17 00:00:00 2001
From: tanaka takashi <isazaka@gmail.com>
Date: Thu, 9 Jul 2026 19:52:03 +0900
Subject: 停止ボタンの不具合修正(停止ボタンを押すと送信できなくなるバグを修正)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

LLMがメッセージ生成途上で画面スクロールしにくい問題を修正
チャットメッセージバルーンをコピペ対応
トークン消費量をリングゲージで可視化(表示タイミングに不具合あり)
---
 src/ollama_web/sessions.py | 43 ++++++++++++++++++++++++++++++++++---------
 1 file changed, 34 insertions(+), 9 deletions(-)

(limited to 'src/ollama_web/sessions.py')

diff --git a/src/ollama_web/sessions.py b/src/ollama_web/sessions.py
index 25f345b..593f2d7 100644
--- a/src/ollama_web/sessions.py
+++ b/src/ollama_web/sessions.py
@@ -89,13 +89,17 @@ class ChatMessage:
     role: str
     content: str
     attachments: list[str] = field(default_factory=list)
+    token_count: dict[str, int] | None = None
 
     def to_dict(self) -> dict[str, Any]:
-        return {
+        out = {
             "role": self.role,
             "content": self.content,
             "attachments": list(self.attachments),
         }
+        if self.token_count is not None:
+            out["token_count"] = dict(self.token_count)
+        return out
 
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> ChatMessage:
@@ -103,6 +107,7 @@ class ChatMessage:
             role=data.get("role", "user"),
             content=data.get("content", ""),
             attachments=list(data.get("attachments", [])),
+            token_count=data.get("token_count"),
         )
 
 
@@ -213,17 +218,18 @@ class SessionStore:
         role: str,
         content: str,
         attachments: list[str] | None = None,
+        token_count: dict[str, int] | None = None,
     ) -> dict[str, Any] | None:
         session = self.get(session_id)
         if session is None:
             return None
-        session["messages"].append(
-            ChatMessage(
-                role=role,
-                content=content,
-                attachments=list(attachments or []),
-            ).to_dict()
-        )
+        msg = ChatMessage(
+            role=role,
+            content=content,
+            attachments=list(attachments or []),
+            token_count=token_count,
+        ).to_dict()
+        session["messages"].append(msg)
         if role == "user" and len(session["messages"]) <= 2:
             first_user = next(
                 (m for m in session["messages"] if m["role"] == "user"),
@@ -232,9 +238,27 @@ class SessionStore:
             if first_user:
                 title = first_user["content"].strip().replace("\n", " ")[:40]
                 session["title"] = title + ("…" if len(first_user["content"]) > 40 else "")
+        self._update_session_token_count(session)
         self.save(session)
         return session
 
+    def _update_session_token_count(self, session: dict[str, Any]) -> None:
+        """Recalculate cumulative token_count for the session."""
+        total_prompt = 0
+        total_eval = 0
+        for m in session.get("messages", []):
+            tc = m.get("token_count") or {}
+            total_prompt += int(tc.get("prompt_eval_count", 0) or 0)
+            total_eval += int(tc.get("eval_count", 0) or 0)
+        if total_prompt or total_eval:
+            session["token_count"] = {
+                "prompt_eval_count": total_prompt,
+                "eval_count": total_eval,
+                "total_count": total_prompt + total_eval,
+            }
+        else:
+            session.pop("token_count", None)
+
     def add_file(
         self,
         session_id: str,
@@ -334,5 +358,6 @@ class SessionStore:
         session["messages"] = []
         session["files"] = []
         session["updated_at"] = _now()
+        session.pop("token_count", None)
         self._save(session_id, session)
-        return session
+        return session
\ No newline at end of file
-- 
cgit v1.2.3