app.py

import streamlit as st
import requests
import json
from duckduckgo_search import DDGS
import re
import clipboard
import io
import os
import pandas as pd
from PyPDF2 import PdfFileReader
from PIL import Image
import pytesseract
from bs4 import BeautifulSoup
import tiktoken
from langchain import LLMChain, PromptTemplate

# OpenRouter APIキーの設定
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY").strip()

# Streamlitアプリケーションの設定
st.set_page_config(page_title="Orb AI", page_icon="🤖", layout="wide")

# セッション状態の初期化
if 'messages' not in st.session_state:
    st.session_state.messages = []
if 'uploaded_file_content' not in st.session_state:
    st.session_state.uploaded_file_content = None

def read_pdf(file):
    pdf = PdfFileReader(file)
    content = ""
    for page_num in range(pdf.getNumPages()):
        page = pdf.getPage(page_num)
        content += page.extract_text()
    return content

def read_csv(file):
    df = pd.read_csv(file)
    return df.to_string()

def read_image(file):
    image = Image.open(file)
    text = pytesseract.image_to_string(image)
    return text

def fetch_webpage_text(url):
    response = requests.get(url)
    if response.status_code != 200:
        return "Web page could not be retrieved"
    soup = BeautifulSoup(response.text, 'html.parser')
    paragraphs = soup.find_all('p')
    text_content = ' '.join([paragraph.get_text() for paragraph in paragraphs])
    return text_content

def truncate_prompt(prompt, max_tokens):
    tokenizer = tiktoken.get_encoding("cl100k_base")
    tokens = tokenizer.encode(prompt)
    if len(tokens) > max_tokens:
        truncated_tokens = tokens[:max_tokens]
        truncated_prompt = tokenizer.decode(truncated_tokens)
        return truncated_prompt
    return prompt

def call_openrouter_api(prompt, model):
    max_context_tokens = 4096
    completion_tokens = 500
    max_prompt_tokens = max_context_tokens - completion_tokens
    truncated_prompt = truncate_prompt(prompt, max_prompt_tokens)
    response = requests.post(
        url="https://openrouter.ai/api/v1/chat/completions",
        headers={
            "Authorization": f"Bearer {OPENROUTER_API_KEY}",
            "HTTP-Referer": "https://supertakerin2-comcomgptfree.hf.space/",
            "X-Title": "Orb AI",
        },
        data=json.dumps({
            "model": model,
            "messages": [{"role": "user", "content": truncated_prompt}]
        })
    )
    if response.status_code == 200:
        try:
            response_json = response.json()
            if "choices" in response_json:
                return response_json["choices"][0]["message"]["content"]
            else:
                st.error("API response does not contain 'choices'")
                return "API response does not contain 'choices'"
        except json.JSONDecodeError:
            st.error("API response JSON decoding failed.")
            return "API response JSON decoding failed."
    else:
        st.error(f"API request failed status code: {response.status_code}, message: {response.text}")
        return "API request failed"

# サイドバーにAIモデルの選択、モード選択、ファイルアップロードを配置
with st.sidebar:
    ai_model = st.selectbox(
        "Select AI model",
        ["openchat/openchat-7b:free", "microsoft/phi-3-mini-128k-instruct:free", "meta-llama/llama-3.1-8b-instruct:free", "google/gemma-7b-it:free", "undi95/toppy-m-7b:free", "mistralai/mistral-7b-instruct:free"],
        key="model_select"
    )
    mode = st.selectbox(
        "Select mode",
        ["focus", "writing", "video", "summary"],
        key="mode_select"
    )
    # ファイルアップロード機能をサイドバーに移動
    uploaded_file = st.file_uploader("Upload file", type=["txt", "pdf", "csv", "py", "html", "css", "js", "cs", "php", "java", "jpg"])
    if uploaded_file is not None:
        file_extension = uploaded_file.name.split('.')[-1].lower()
        if file_extension in ['txt', 'py', 'html', 'css', 'js', 'cs', 'php', 'java']:
            stringio = io.StringIO(uploaded_file.getvalue().decode("utf-8"))
            st.session_state.uploaded_file_content = stringio.read()
        elif file_extension == 'pdf':
            st.session_state.uploaded_file_content = read_pdf(uploaded_file)
        elif file_extension == 'csv':
            st.session_state.uploaded_file_content = read_csv(uploaded_file)
        elif file_extension in ['jpg']:
            st.session_state.uploaded_file_content = read_image(uploaded_file)
        st.success(f"file '{uploaded_file.name}' was uploaded")

    # 会話をクリアするボタン
    if st.button("🔥 Delete conversation"):
        st.session_state.messages = []
        st.session_state.uploaded_file_content = None

# メインコンテンツ
st.title("Orb AI")

# チャット履歴の表示
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# ユーザー入力
if prompt := st.chat_input("Please enter your message."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # アップロードされたファイルの内容を考慮してプロンプトを作成
    if st.session_state.uploaded_file_content:
        prompt = f"以下のファイル内容を参考にして回答してください:\n\n{st.session_state.uploaded_file_content}\n\nユーザーの質問: {prompt}"

    full_response = ""

    if mode == "focus":
        # DuckDuckGo検索の実行
        with DDGS() as ddgs:
            search_results = list(ddgs.text(prompt, max_results=4))
            search_results_text = "\n".join([result["body"] for result in search_results])
            urls = "\n".join([result["href"] for result in search_results[:4]])
            prompt = f"質問: {prompt}\n\n検索結果: {search_results_text}\n\n上記の情報に基づいて、詳細な回答を提供してください："
            full_response = call_openrouter_api(prompt, ai_model)
            # 回答の下に検索結果のURLを表示
            with st.chat_message("assistant"):
                st.markdown(full_response)
                st.markdown("### Sources:")
                st.markdown(urls)
            
    elif mode == "writing":
        prompt = f"次のトピックについて包括的な記事を書いてください: {prompt}\n\n導入、主要なポイント、および結論を含めてください。"
        full_response = call_openrouter_api(prompt, ai_model)
        with st.chat_message("assistant"):
            st.markdown(full_response)
        
    elif mode == "summary":
        prompt = f"次のテキストを簡潔に要約してください:\n\n{prompt}\n\n主要なポイントを捉えた簡潔な要約を提供してください。"
        full_response = call_openrouter_api(prompt, ai_model)
        with st.chat_message("assistant"):
            st.markdown(full_response)
        
    elif mode == "video":
        # DuckDuckGo検索を使用してYouTubeの動画を検索
        with DDGS() as ddgs:
            search_results = list(ddgs.text(f"{prompt} site:youtube.com", max_results=5))
            if search_results:
                for result in search_results[:4]:
                    video_url = result['href']
                    video_id = None
                    youtube_patterns = [
                        r"(?<=v=)[^/]+",
                        r"(?<=be/)[^/]+",
                        r"(?<=embed/)[^/]+"
                    ]
                    for pattern in youtube_patterns:
                        match = re.search(pattern, video_url)
                        if match:
                            video_id = match.group()
                            break
                    if video_id:
                        st.markdown(f'<iframe width="560" height="315" src="https://www.youtube.com/embed/{video_id}" frameborder="0" allowfullscreen></iframe>', unsafe_allow_html=True)
                        st.markdown(f"Related Videos: {result['title']}")
                    else:
                        st.warning(f"Video ID not found: {video_url}")
            else:
                full_response = "No related videos were found."
                with st.chat_message("assistant"):
                    st.markdown(full_response)

    # コピー機能の追加
    if st.button("📋 Copy the generated text"):
        clipboard.copy(full_response)
        st.success("Text has been copied to the clipboard!")