personal memory agent
Merge pull request #126 from kognova/codex/review-chat-view-and-server-access
Enable STDIO MCP server and Gemini chat integration
authored by
Jer Miller
and committed by
GitHub
10563e7a
b6a64752
+85
-14
+29
-9
dream/views/chat.py
+29
-9
dream/views/chat.py
···
1
1
from __future__ import annotations
2
2
3
+
import asyncio
3
4
import os
4
-
from typing import Any
5
+
from typing import Any, List
5
6
6
7
from flask import Blueprint, jsonify, render_template, request
7
8
from google import genai
8
9
from google.genai import types
9
10
11
+
from mcp_tools import sunstone_toolset
10
12
from think.models import GEMINI_FLASH
11
13
12
14
bp = Blueprint("chat", __name__, template_folder="../templates")
13
15
16
+
loop = asyncio.new_event_loop()
17
+
_toolset = None
18
+
19
+
20
+
async def _get_toolset():
21
+
global _toolset
22
+
if _toolset is None:
23
+
_toolset = await sunstone_toolset()
24
+
return _toolset
25
+
26
+
27
+
def ask_gemini(prompt: str, attachments: List[str], api_key: str) -> str:
28
+
genai.configure(api_key=api_key)
29
+
toolset = loop.run_until_complete(_get_toolset())
30
+
model = genai.GenerativeModel(
31
+
model_name=GEMINI_FLASH,
32
+
tools=[toolset],
33
+
tool_config=types.ToolConfig(
34
+
function_calling_config=types.FunctionCallingConfig(mode="AUTO")
35
+
),
36
+
)
37
+
chat = model.start_chat(history=[])
38
+
return chat.send_message([prompt] + attachments).text
39
+
14
40
15
41
@bp.route("/chat")
16
42
def chat_page() -> str:
···
26
52
payload = request.get_json(force=True)
27
53
message = payload.get("message", "")
28
54
attachments = payload.get("attachments", [])
29
-
contents = [message] + attachments
30
55
31
-
client = genai.Client(api_key=api_key)
32
-
response = client.models.generate_content(
33
-
model=GEMINI_FLASH,
34
-
contents=contents,
35
-
config=types.GenerateContentConfig(temperature=0.3, max_output_tokens=1024),
36
-
)
37
-
return jsonify({"text": response.text})
56
+
answer = loop.run_in_executor(None, ask_gemini, message, attachments, api_key).result()
57
+
return jsonify(text=answer)
+29
mcp_tools.py
+29
mcp_tools.py
···
1
+
import asyncio
2
+
import os
3
+
4
+
_client = None
5
+
6
+
7
+
_SERVER_PATH = os.path.join(os.path.dirname(__file__), "sunstone_server.py")
8
+
_SERVER_URL = os.getenv("SUNSTONE_MCP_URL")
9
+
10
+
11
+
def _get_client():
12
+
global _client
13
+
if _client is None:
14
+
from fastmcp import Client
15
+
16
+
_client = Client(_SERVER_URL or _SERVER_PATH, keep_alive=True)
17
+
return _client
18
+
19
+
20
+
async def sunstone_toolset():
21
+
"""Return a Tool object ready for Gemini/Chat."""
22
+
client = _get_client()
23
+
await client.initialize()
24
+
return client.as_toolset(label="Sunstone")
25
+
26
+
27
+
async def close_client():
28
+
if _client is not None:
29
+
await _client.close()
+13
sunstone_server.py
+13
sunstone_server.py
···
1
+
import argparse
2
+
3
+
from dotenv import load_dotenv
4
+
5
+
from think.mcp_server import create_server
6
+
7
+
if __name__ == "__main__":
8
+
load_dotenv()
9
+
ap = argparse.ArgumentParser()
10
+
ap.add_argument("--journal", required=True, help="Path to the .jrnl database")
11
+
opts = ap.parse_args()
12
+
13
+
create_server(opts.journal).run()
+14
-5
think/mcp_server.py
+14
-5
think/mcp_server.py
···
46
46
load_dotenv()
47
47
parser = argparse.ArgumentParser()
48
48
parser.add_argument("--port", type=int, default=8000)
49
+
parser.add_argument(
50
+
"--stdio",
51
+
action="store_true",
52
+
help="Run using STDIO transport instead of HTTP",
53
+
)
49
54
args = parser.parse_args()
50
55
51
56
journal = os.getenv("JOURNAL_PATH") or parser.error("JOURNAL_PATH not set")
52
-
create_server(journal).run(
53
-
"streamable-http",
54
-
host="0.0.0.0",
55
-
port=args.port,
56
-
)
57
+
server = create_server(journal)
58
+
if args.stdio:
59
+
server.run()
60
+
else:
61
+
server.run(
62
+
"streamable-http",
63
+
host="0.0.0.0",
64
+
port=args.port,
65
+
)
57
66
58
67
59
68
if __name__ == "__main__":